Spaces:
Runtime error
Runtime error
| import spaces | |
| import json | |
| import math | |
| import os | |
| import traceback | |
| from io import BytesIO | |
| from typing import Any, Dict, List, Optional, Tuple, Iterable | |
| import re | |
| import time | |
| from threading import Thread | |
| from io import BytesIO | |
| import uuid | |
| import tempfile | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import supervision as sv | |
| from transformers import ( | |
| Qwen2_5_VLForConditionalGeneration, | |
| Glm4vForConditionalGeneration, | |
| Qwen2VLForConditionalGeneration, | |
| AutoModelForCausalLM, | |
| AutoProcessor, | |
| TextIteratorStreamer, | |
| ) | |
| from gradio.themes import Soft | |
| from gradio.themes.utils import colors, fonts, sizes | |
| # --- Theme Definition --- | |
| # Define a new color palette for Blue | |
| colors.blue_theme_color = colors.Color( | |
| name="blue_theme_color", | |
| c50="#E6E6FF", | |
| c100="#CCCCFF", | |
| c200="#9999FF", | |
| c300="#6666FF", | |
| c400="#3333FF", | |
| c500="#0000FF", # Base Blue color | |
| c600="#0000D9", | |
| c700="#0000B3", | |
| c800="#000080", | |
| c900="#000066", | |
| c950="#000033", | |
| ) | |
| class BlueTheme(Soft): | |
| def __init__( | |
| self, | |
| *, | |
| primary_hue: colors.Color | str = colors.gray, | |
| secondary_hue: colors.Color | str = colors.blue_theme_color, | |
| neutral_hue: colors.Color | str = colors.slate, | |
| text_size: sizes.Size | str = sizes.text_lg, | |
| font: fonts.Font | str | Iterable[fonts.Font | str] = ( | |
| fonts.GoogleFont("Outfit"), "Arial", "sans-serif", | |
| ), | |
| font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( | |
| fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", | |
| ), | |
| ): | |
| super().__init__( | |
| primary_hue=primary_hue, | |
| secondary_hue=secondary_hue, | |
| neutral_hue=neutral_hue, | |
| text_size=text_size, | |
| font=font, | |
| font_mono=font_mono, | |
| ) | |
| super().set( | |
| background_fill_primary="*primary_50", | |
| background_fill_primary_dark="*primary_900", | |
| body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", | |
| body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", | |
| button_primary_text_color="white", | |
| button_primary_text_color_hover="white", | |
| button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", | |
| button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", | |
| button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", | |
| button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", | |
| slider_color="*secondary_500", | |
| slider_color_dark="*secondary_600", | |
| block_title_text_weight="600", | |
| block_border_width="2px", | |
| block_shadow="*shadow_drop_lg", | |
| button_primary_shadow="*shadow_drop_lg", | |
| button_large_padding="12px", | |
| block_label_background_fill="*primary_200", | |
| ) | |
| # Instantiate the theme | |
| blue_theme = BlueTheme() | |
| # --- Constants and Model Setup --- | |
| MAX_INPUT_TOKEN_LENGTH = 4096 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("--- System Information ---") | |
| print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES")) | |
| print("torch.__version__ =", torch.__version__) | |
| print("torch.version.cuda =", torch.version.cuda) | |
| print("CUDA available:", torch.cuda.is_available()) | |
| print("CUDA device count:", torch.cuda.device_count()) | |
| if torch.cuda.is_available(): | |
| print("Current device:", torch.cuda.current_device()) | |
| print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device())) | |
| print("Using device:", device) | |
| print("--------------------------") | |
| # --- Model Loading --- | |
| # Load Camel-Doc-OCR-062825 | |
| print("Loading Camel-Doc-OCR-062825...") | |
| MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825" | |
| processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True) | |
| model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_M, | |
| trust_remote_code=True, | |
| torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| print("Camel-Doc-OCR-062825 loaded.") | |
| # GLM-4.1V-9B-Thinking | |
| print("Loading GLM-4.1V-9B-Thinking") | |
| MODEL_ID_T = "zai-org/GLM-4.1V-9B-Thinking" | |
| processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True) | |
| model_t = Glm4vForConditionalGeneration.from_pretrained( | |
| MODEL_ID_T, | |
| trust_remote_code=True, | |
| torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| print("GLM-4.1V-9B-Thinking loaded.") | |
| # Load moondream3 | |
| print("Loading moondream3-preview...") | |
| MODEL_ID_MD3 = "moondream/moondream3-preview" | |
| model_md3 = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID_MD3, | |
| trust_remote_code=True, | |
| torch_dtype=torch.bfloat16, | |
| device_map={"": "cuda"}, | |
| ) | |
| model_md3.compile() | |
| print("moondream3-preview loaded and compiled.") | |
| # --- Moondream3 Utility Functions --- | |
| def create_annotated_image(image, detection_result, object_name="Object"): | |
| if not isinstance(detection_result, dict) or "objects" not in detection_result: | |
| return image | |
| original_width, original_height = image.size | |
| annotated_image = np.array(image.convert("RGB")) | |
| bboxes = [] | |
| labels = [] | |
| for i, obj in enumerate(detection_result["objects"]): | |
| x_min = int(obj["x_min"] * original_width) | |
| y_min = int(obj["y_min"] * original_height) | |
| x_max = int(obj["x_max"] * original_width) | |
| y_max = int(obj["y_max"] * original_height) | |
| x_min = max(0, min(x_min, original_width)) | |
| y_min = max(0, min(y_min, original_height)) | |
| x_max = max(0, min(x_max, original_width)) | |
| y_max = max(0, min(y_max, original_height)) | |
| if x_max > x_min and y_max > y_min: | |
| bboxes.append([x_min, y_min, x_max, y_max]) | |
| labels.append(f"{object_name} {i+1}") | |
| if not bboxes: | |
| return image | |
| detections = sv.Detections( | |
| xyxy=np.array(bboxes, dtype=np.float32), | |
| class_id=np.arange(len(bboxes)) | |
| ) | |
| bounding_box_annotator = sv.BoxAnnotator( | |
| thickness=3, | |
| color_lookup=sv.ColorLookup.INDEX | |
| ) | |
| label_annotator = sv.LabelAnnotator( | |
| text_thickness=2, | |
| text_scale=0.6, | |
| color_lookup=sv.ColorLookup.INDEX | |
| ) | |
| annotated_image = bounding_box_annotator.annotate( | |
| scene=annotated_image, detections=detections | |
| ) | |
| annotated_image = label_annotator.annotate( | |
| scene=annotated_image, detections=detections, labels=labels | |
| ) | |
| return Image.fromarray(annotated_image) | |
| def create_point_annotated_image(image, point_result): | |
| if not isinstance(point_result, dict) or "points" not in point_result: | |
| return image | |
| original_width, original_height = image.size | |
| annotated_image = np.array(image.convert("RGB")) | |
| points = [] | |
| for point in point_result["points"]: | |
| x = int(point["x"] * original_width) | |
| y = int(point["y"] * original_height) | |
| points.append([x, y]) | |
| if points: | |
| points_array = np.array(points).reshape(1, -1, 2) | |
| key_points = sv.KeyPoints(xy=points_array) | |
| vertex_annotator = sv.VertexAnnotator(radius=8, color=sv.Color.RED) | |
| annotated_image = vertex_annotator.annotate( | |
| scene=annotated_image, key_points=key_points | |
| ) | |
| return Image.fromarray(annotated_image) | |
| def detect_objects_md3(image, prompt, task_type, max_objects): | |
| STANDARD_SIZE = (1024, 1024) | |
| if image is None: | |
| raise gr.Error("Please upload an image.") | |
| image.thumbnail(STANDARD_SIZE) | |
| t0 = time.perf_counter() | |
| if task_type == "Object Detection": | |
| settings = {"max_objects": max_objects} if max_objects > 0 else {} | |
| result = model_md3.detect(image, prompt, settings=settings) | |
| annotated_image = create_annotated_image(image, result, prompt) | |
| elif task_type == "Point Detection": | |
| result = model_md3.point(image, prompt) | |
| annotated_image = create_point_annotated_image(image, result) | |
| elif task_type == "Caption": | |
| result = model_md3.caption(image, length="normal") | |
| annotated_image = image | |
| else: | |
| result = model_md3.query(image=image, question=prompt, reasoning=True) | |
| annotated_image = image | |
| elapsed_ms = (time.perf_counter() - t0) * 1_000 | |
| if isinstance(result, dict): | |
| if "objects" in result: | |
| output_text = f"Found {len(result['objects'])} objects:\n" | |
| for i, obj in enumerate(result['objects'], 1): | |
| output_text += f"\n{i}. Bounding box: ({obj['x_min']:.3f}, {obj['y_min']:.3f}, {obj['x_max']:.3f}, {obj['y_max']:.3f})" | |
| elif "points" in result: | |
| output_text = f"Found {len(result['points'])} points:\n" | |
| for i, point in enumerate(result['points'], 1): | |
| output_text += f"\n{i}. Point: ({point['x']:.3f}, {point['y']:.3f})" | |
| elif "caption" in result: | |
| output_text = result['caption'] | |
| elif "answer" in result: | |
| output_text = f"Reasoning: {result.get('reasoning', 'N/A')}\n\nAnswer: {result['answer']}" | |
| else: | |
| output_text = json.dumps(result, indent=2) | |
| else: | |
| output_text = str(result) | |
| timing_text = f"Inference time: {elapsed_ms:.0f} ms" | |
| return annotated_image, output_text, timing_text | |
| # --- Core Application Logic (for other models) --- | |
| def process_document_stream( | |
| model_name: str, | |
| image: Image.Image, | |
| prompt_input: str, | |
| max_new_tokens: int, | |
| temperature: float, | |
| top_p: float, | |
| top_k: int, | |
| repetition_penalty: float | |
| ): | |
| """ | |
| Main generator function for models other than Moondream3. | |
| """ | |
| if image is None: | |
| yield "Please upload an image." | |
| return | |
| if not prompt_input or not prompt_input.strip(): | |
| yield "Please enter a prompt." | |
| return | |
| # Select processor and model based on dropdown choice | |
| if model_name == "Camel-Doc-OCR-062825 (OCR)": | |
| processor, model = processor_m, model_m | |
| elif model_name == "GLM-4.1V-9B (Thinking)": | |
| processor, model = processor_t, model_t | |
| else: | |
| yield "Invalid model selected." | |
| return | |
| messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_input}]}] | |
| prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device) | |
| streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True) | |
| generation_kwargs = { | |
| **inputs, | |
| "streamer": streamer, | |
| "max_new_tokens": max_new_tokens, | |
| "temperature": temperature, | |
| "top_p": top_p, | |
| "top_k": top_k, | |
| "repetition_penalty": repetition_penalty, | |
| "do_sample": True if temperature > 0 else False | |
| } | |
| thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| buffer = "" | |
| for new_text in streamer: | |
| buffer += new_text | |
| # Clean up potential model-specific tokens | |
| buffer = buffer.replace("<|im_end|>", "").replace("</s>", "") | |
| time.sleep(0.01) | |
| yield buffer | |
| def create_gradio_interface(): | |
| """Builds and returns the Gradio web interface.""" | |
| css = """ | |
| .main-container { max-width: 1400px; margin: 0 auto; } | |
| #main-title h1 {font-size: 2.2em !important;} | |
| """ | |
| with gr.Blocks(theme=blue_theme, css=css) as demo: | |
| gr.Markdown("# **Multimodal VLM v1.0**", elem_id="main-title") | |
| gr.Markdown("Explore the capabilities of various Vision Language Models for tasks like OCR, VQA, and Object Detection.") | |
| with gr.Tabs(): | |
| # --- TAB 1: Document and General VLMs --- | |
| with gr.TabItem("📄 Document & General VLM"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| #gr.Markdown("### 1. Configure Inputs") | |
| model_choice = gr.Dropdown( | |
| choices=["Camel-Doc-OCR-062825 (OCR)", "GLM-4.1V-9B (Thinking)"], | |
| label="Select Model", value= "Camel-Doc-OCR-062825 (OCR)" | |
| ) | |
| image_input_doc = gr.Image(label="Upload Image", type="pil", sources=['upload'], height=280) | |
| prompt_input_doc = gr.Textbox(label="Query Input", placeholder="e.g., 'Transcribe the text in this document.'") | |
| with gr.Accordion("Advanced Generation Settings", open=False): | |
| max_new_tokens = gr.Slider(minimum=256, maximum=4096, value=2048, step=128, label="Max New Tokens") | |
| temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7) | |
| top_p = gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, step=0.05, value=0.9) | |
| top_k = gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=40) | |
| repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1) | |
| process_btn = gr.Button("Submit", variant="primary") | |
| clear_btn = gr.Button("Clear", variant="secondary") | |
| with gr.Column(scale=2): | |
| #gr.Markdown("### 2. View Output") | |
| with gr.Tab("Output Stream"): | |
| output_stream = gr.Textbox(label="Model Output", interactive=False, lines=24, show_copy_button=True) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/1.jpg", "Transcribe this receipt."], | |
| ["examples/2.jpg", "Extract the content."], | |
| ["examples/3.jpg", "OCR the image."], | |
| ], | |
| inputs=[image_input_doc, prompt_input_doc] | |
| ) | |
| # --- TAB 2: Moondream3 Lab --- | |
| with gr.TabItem("🌝 Moondream3"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| md3_image_input = gr.Image(label="Upload an image", type="pil", height=400) | |
| md3_task_type = gr.Radio( | |
| choices=["Object Detection", "Point Detection", "Caption", "Visual Question Answering"], | |
| label="Task Type", value="Object Detection" | |
| ) | |
| md3_prompt_input = gr.Textbox( | |
| label="Prompt (object to detect/question to ask)", | |
| placeholder="e.g., 'car', 'person', 'What's in this image?'" | |
| ) | |
| md3_max_objects = gr.Number( | |
| label="Max Objects (for Object Detection only)", | |
| value=10, minimum=1, maximum=50, step=1, visible=True | |
| ) | |
| md3_generate_btn = gr.Button(value="Submit", variant="primary") | |
| with gr.Column(scale=1): | |
| md3_output_image = gr.Image(type="pil", label="Result", height=400) | |
| md3_output_textbox = gr.Textbox(label="Model Response", lines=10, show_copy_button=True) | |
| md3_output_time = gr.Markdown() | |
| gr.Examples( | |
| examples=[ | |
| ["md3/1.jpg", "Object Detection", "boats", 7], | |
| ["md3/2.jpg", "Point Detection", "children", 7], | |
| ["md3/3.png", "Caption", "", 5], | |
| ["md3/4.jpeg", "Visual Question Answering", "Analyze the GDP trend over the years.", 5], | |
| ], | |
| inputs=[md3_image_input, md3_task_type, md3_prompt_input, md3_max_objects], | |
| label="Click an example to populate inputs" | |
| ) | |
| process_btn.click( | |
| fn=process_document_stream, | |
| inputs=[model_choice, image_input_doc, prompt_input_doc, max_new_tokens, temperature, top_p, top_k, repetition_penalty], | |
| outputs=[output_stream] | |
| ) | |
| clear_btn.click(lambda: (None, "", ""), outputs=[image_input_doc, prompt_input_doc, output_stream]) | |
| # Moondream3 Tab | |
| def update_max_objects_visibility(task): | |
| return gr.update(visible=(task == "Object Detection")) | |
| md3_task_type.change(fn=update_max_objects_visibility, inputs=[md3_task_type], outputs=[md3_max_objects]) | |
| md3_generate_btn.click( | |
| fn=detect_objects_md3, | |
| inputs=[md3_image_input, md3_prompt_input, md3_task_type, md3_max_objects], | |
| outputs=[md3_output_image, md3_output_textbox, md3_output_time] | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = create_gradio_interface() | |
| demo.queue(max_size=50).launch(ssr_mode=False, mcp_server=True, show_error=True) |