Update app.py

app.py

@@ -20,6 +20,8 @@ from transformers import (
     AutoModelForImageTextToText,
     AutoProcessor,
     TextIteratorStreamer,
+    AutoModel,
+    AutoTokenizer,
 )
 from transformers.image_utils import load_image
 
@@ -30,6 +32,8 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+# --- Original Models ---
+
 # Load DREX-062225-exp
 MODEL_ID_X = "prithivMLmods/DREX-062225-exp"
 processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
@@ -68,14 +72,16 @@ model_j = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load LMM-R1-MGT-PerceReason
-MODEL_ID_F = "VLM-Reasoner/LMM-R1-MGT-PerceReason"
-processor_f = AutoProcessor.from_pretrained(MODEL_ID_F, trust_remote_code=True)
-model_f = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_F,
+# --- Load New Model: openbmb/MiniCPM-V-4 ---
+MODEL_ID_V4 = 'openbmb/MiniCPM-V-4'
+model_v4 = AutoModel.from_pretrained(
+    MODEL_ID_V4,
     trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
+    torch_dtype=torch.bfloat16,
+    attn_implementation='sdpa'
+).eval().to(device)
+tokenizer_v4 = AutoTokenizer.from_pretrained(MODEL_ID_V4, trust_remote_code=True)
+
 
 def downsample_video(video_path):
     """
@@ -86,7 +92,8 @@ def downsample_video(video_path):
     total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
     fps = vidcap.get(cv2.CAP_PROP_FPS)
     frames = []
-    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    # Use a maximum of 10 frames to avoid excessive memory usage
+    frame_indices = np.linspace(0, total_frames - 1, min(total_frames, 10), dtype=int)
     for i in frame_indices:
         vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
         success, image = vidcap.read()
@@ -108,44 +115,46 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     """
     Generates responses using the selected model for image input.
     """
+    if image is None:
+        yield "Please upload an image.", "Please upload an image."
+        return
+
+    # Handle the new model separately due to its different API
+    if model_name == "openbmb/MiniCPM-V-4":
+        msgs = [{'role': 'user', 'content': [image, text]}]
+        try:
+            answer = model_v4.chat(
+                image=image.convert('RGB'),
+                msgs=msgs,
+                tokenizer=tokenizer_v4,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+            )
+            yield answer, answer
+        except Exception as e:
+            yield f"Error: {e}", f"Error: {e}"
+        return
+
+    # Original model selection logic
     if model_name == "DREX-062225-7B-exp":
-        processor = processor_x
-        model = model_x
+        processor, model = processor_x, model_x
     elif model_name == "olmOCR-7B-0225-preview":
-        processor = processor_o
-        model = model_o
+        processor, model = processor_o, model_o
     elif model_name == "Typhoon-OCR-3B":
-        processor = processor_t
-        model = model_t
+        processor, model = processor_t, model_t
     elif model_name == "Lumian-VLR-7B-Thinking":
-        processor = processor_j
-        model = model_j
-    elif model_name == "LMM-R1-MGT-PerceReason":
-        processor = processor_f
-        model = model_f
+        processor, model = processor_j, model_j
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
 
-    if image is None:
-        yield "Please upload an image.", "Please upload an image."
-        return
-
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": text},
-        ]
-    }]
+    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}]
     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=False,
-        max_length=MAX_INPUT_TOKEN_LENGTH
+        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}
@@ -167,57 +176,64 @@ def generate_video(model_name: str, text: str, video_path: str,
     """
     Generates responses using the selected model for video input.
     """
+    if video_path is None:
+        yield "Please upload a video.", "Please upload a video."
+        return
+
+    frames_with_ts = downsample_video(video_path)
+    if not frames_with_ts:
+        yield "Could not process video.", "Could not process video."
+        return
+
+    # Handle the new model separately
+    if model_name == "openbmb/MiniCPM-V-4":
+        images = [frame for frame, ts in frames_with_ts]
+        content = [text] + images
+        msgs = [{'role': 'user', 'content': content}]
+        try:
+            answer = model_v4.chat(
+                image=images[0].convert('RGB'),
+                msgs=msgs,
+                tokenizer=tokenizer_v4,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+            )
+            yield answer, answer
+        except Exception as e:
+            yield f"Error: {e}", f"Error: {e}"
+        return
+
+    # Original model selection logic
     if model_name == "DREX-062225-7B-exp":
-        processor = processor_x
-        model = model_x
+        processor, model = processor_x, model_x
     elif model_name == "olmOCR-7B-0225-preview":
-        processor = processor_o
-        model = model_o
+        processor, model = processor_o, model_o
     elif model_name == "Typhoon-OCR-3B":
-        processor = processor_t
-        model = model_t
+        processor, model = processor_t, model_t
     elif model_name == "Lumian-VLR-7B-Thinking":
-        processor = processor_j
-        model = model_j
-    elif model_name == "LMM-R1-MGT-PerceReason":
-        processor = processor_f
-        model = model_f
+        processor, model = processor_j, model_j
     else:
         yield "Invalid model selected.", "Invalid model selected."
         return
 
-    if video_path is None:
-        yield "Please upload a video.", "Please upload a video."
-        return
+    # Prepare messages for Qwen-style models
+    messages = [{"role": "user", "content": [{"type": "text", "text": text}]}]
+    for frame, timestamp in frames_with_ts:
+        messages[0]["content"].append({"type": "image", "image": frame})
 
-    frames = downsample_video(video_path)
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-        {"role": "user", "content": [{"type": "text", "text": text}]}
-    ]
-    for frame in frames:
-        image, timestamp = frame
-        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
-        messages[1]["content"].append({"type": "image", "image": image})
-    inputs = processor.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        truncation=False,
-        max_length=MAX_INPUT_TOKEN_LENGTH
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    images_for_processor = [frame for frame, ts in frames_with_ts]
+    inputs = processor(
+        text=[prompt_full], images=images_for_processor, 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,
-        "do_sample": True,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "repetition_penalty": repetition_penalty,
+        **inputs, "streamer": streamer, "max_new_tokens": max_new_tokens,
+        "do_sample": True, "temperature": temperature, "top_p": top_p,
+        "top_k": top_k, "repetition_penalty": repetition_penalty,
     }
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
@@ -228,14 +244,6 @@ def generate_video(model_name: str, text: str, video_path: str,
         time.sleep(0.01)
         yield buffer, buffer
 
-def save_to_md(output_text):
-    """
-    Saves the output text to a Markdown file and returns the file path for download.
-    """
-    file_path = f"result_{uuid.uuid4()}.md"
-    with open(file_path, "w") as f:
-        f.write(output_text)
-    return file_path
 
 # Define examples for image and video inference
 image_examples = [
@@ -254,18 +262,9 @@ video_examples = [
 
 # Added CSS to style the output area as a "Canvas"
 css = """
-.submit-btn {
-    background-color: #2980b9 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #3498db !important;
-}
-.canvas-output {
-    border: 2px solid #4682B4;
-    border-radius: 10px;
-    padding: 20px;
-}
+.submit-btn { background-color: #2980b9 !important; color: white !important; }
+.submit-btn:hover { background-color: #3498db !important; }
+.canvas-output { border: 2px solid #4682B4; border-radius: 10px; padding: 20px; }
 """
 
 # Create the Gradio Interface
@@ -278,19 +277,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
                     image_upload = gr.Image(type="pil", label="Image")
                     image_submit = gr.Button("Submit", elem_classes="submit-btn")
-                    gr.Examples(
-                        examples=image_examples,
-                        inputs=[image_query, image_upload]
-                    )
+                    gr.Examples(examples=image_examples, inputs=[image_query, image_upload])
                 with gr.TabItem("Video Inference"):
                     video_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
                     video_upload = gr.Video(label="Video")
                     video_submit = gr.Button("Submit", elem_classes="submit-btn")
-                    gr.Examples(
-                        examples=video_examples,
-                        inputs=[video_query, video_upload]
-                    )
-                    
+                    gr.Examples(examples=video_examples, inputs=[video_query, video_upload])
+
             with gr.Accordion("Advanced options", open=False):
                 max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
                 temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
@@ -302,19 +295,21 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             with gr.Column(elem_classes="canvas-output"):
                 gr.Markdown("## Output")
                 output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=2, show_copy_button=True)
-                with gr.Accordion("(Result.md)", open=False):                
+                with gr.Accordion("(Result.md)", open=False):
                     markdown_output = gr.Markdown(label="(Result.Md)")
             model_choice = gr.Radio(
-                choices=["Lumian-VLR-7B-Thinking", "DREX-062225-7B-exp", "olmOCR-7B-0225-preview", "LMM-R1-MGT-PerceReason", "Typhoon-OCR-3B"],
+                choices=[ "openbmb/MiniCPM-V-4", "Lumian-VLR-7B-Thinking", "DREX-062225-7B-exp", "olmOCR-7B-0225-preview", "Typhoon-OCR-3B"],
                 label="Select Model",
                 value="Lumian-VLR-7B-Thinking"
             )
-            gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-VLM-Thinking/discussions)")          
-            gr.Markdown("> Lumian-VLR-7B-Thinking is a high-fidelity vision-language reasoning model built on Qwen2.5-VL-7B-Instruct, designed for fine-grained multimodal understanding, enhancing image captioning, video reasoning, and document comprehension through explicit grounded reasoning. It is trained first via supervised fine-tuning (SFT) on visually-grounded reasoning traces and then further refined using GRPO reinforcement learning to boost reasoning accuracy.")
-            gr.Markdown("> LMM-R1-MGT-PerceReason is a vision-language model focused on advanced reasoning using a multimodal tree search approach enabling progressive visual-textual slow thinking, improving complex spatial and logical reasoning without fine-tuning. OLMOCR-7B-0225-preview is a 7B parameter open large model designed for OCR tasks with robust text extraction, especially in complex document layouts. ")
-            gr.Markdown("> Typhoon-ocr-3b is a 3B parameter OCR model optimized for efficient and accurate optical character recognition in challenging conditions. DREX-062225-exp is an experimental multimodal model emphasizing strong document reading and extraction capabilities combined with vision-language understanding to support detailed document parsing and reasoning tasks.")
-            gr.Markdown("> ⚠️ Note: Models in this space may not perform well on video inference tasks.")
-            
+            gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-VLM-Thinking/discussions)")
+            gr.Markdown("> **MiniCPM-V 4** is a powerful open-source multimodal model capable of handling various image and text-based tasks with high accuracy.")
+            gr.Markdown("> **Lumian-VLR-7B-Thinking** is a high-fidelity vision-language reasoning model built on Qwen2.5-VL-7B-Instruct, designed for fine-grained multimodal understanding, enhancing image captioning, video reasoning, and document comprehension through explicit grounded reasoning.")
+            gr.Markdown("> **olmOCR-7B-0225-preview** is a 7B parameter open large model designed for OCR tasks with robust text extraction, especially in complex document layouts. ")
+            gr.Markdown("> **Typhoon-ocr-3b** is a 3B parameter OCR model optimized for efficient and accurate optical character recognition in challenging conditions.")
+            gr.Markdown("> **DREX-062225-exp** is an experimental multimodal model emphasizing strong document reading and extraction capabilities combined with vision-language understanding to support detailed document parsing and reasoning tasks.")
+            gr.Markdown("> ⚠️ Note: Video inference performance can vary significantly between models.")
+
     image_submit.click(
         fn=generate_image,
         inputs=[model_choice, image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
@@ -327,4 +322,4 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)
+    demo.queue(max_size=30).launch(share=True, show_error=True)