Update app.py

app.py

@@ -16,8 +16,7 @@ import cv2
 
 from transformers import (
     Qwen2_5_VLForConditionalGeneration,
-    Qwen2VLForConditionalGeneration,
-    AutoModel,
+    Qwen2VLForConditionalGeneration, 
     AutoProcessor,
     AutoTokenizer,
     TextIteratorStreamer,
@@ -57,15 +56,13 @@ model_k = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load llama-nemoretriever-colembed-1b-v1
-MODEL_ID_Y = "nvidia/llama-nemoretriever-colembed-1b-v1"
-#processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
-model_y = AutoModel.from_pretrained(
+# Load remyxai/SpaceOm
+MODEL_ID_Y = "remyxai/SpaceOm"
+processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
+model_y = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_Y,
     trust_remote_code=True,
-    torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2",
-    revision='1f0fdea7f5b19532a750be109b19072d719b8177'
+    torch_dtype=torch.float16
 ).to(device).eval()
 
 
@@ -109,8 +106,8 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "coreOCR-7B-050325-preview":
         processor = processor_k
         model = model_k
-    elif model_name == "llama-nemoretriever-colembed-1b-v1":
-        #processor = processor_y
+    elif model_name == "SpaceOm-3B":
+        processor = processor_y
         model = model_y
     else:
         yield "Invalid model selected."
@@ -166,8 +163,8 @@ def generate_video(model_name: str, text: str, video_path: str,
     elif model_name == "coreOCR-7B-050325-preview":
         processor = processor_k
         model = model_k
-    elif model_name == "llama-nemoretriever-colembed-1b-v1":
-        #processor = processor_y
+    elif model_name == "SpaceOm-3B":
+        processor = processor_y
         model = model_y
     else:
         yield "Invalid model selected."
@@ -270,7 +267,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         with gr.Column():
             output = gr.Textbox(label="Output", interactive=False, lines=2, scale=2)
             model_choice = gr.Radio(
-                choices=["SkyCaptioner-V1", "SpaceThinker-3B", "coreOCR-7B-050325-preview", "llama-nemoretriever-colembed-1b-v1"],
+                choices=["SkyCaptioner-V1", "SpaceThinker-3B", "coreOCR-7B-050325-preview", "SpaceOm-3B"],
                 label="Select Model",
                 value="SkyCaptioner-V1"
             )
@@ -278,7 +275,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/VisionScope-R2/discussions)")
             gr.Markdown("> [SkyCaptioner-V1](https://huggingface.co/Skywork/SkyCaptioner-V1):  structural video captioning model designed to generate high-quality, structural descriptions for video data. It integrates specialized sub-expert models.")
             gr.Markdown("> [SpaceThinker-Qwen2.5VL-3B](https://huggingface.co/remyxai/SpaceThinker-Qwen2.5VL-3B): thinking/reasoning multimodal/vision-language model (VLM) trained to enhance spatial reasoning.")
-            gr.Markdown("> [coreOCR-7B-050325-preview](https://huggingface.co/prithivMLmods/coreOCR-7B-050325-preview): model is a fine-tuned version of qwen/qwen2-vl-7b, optimized for document-level optical character recognition (ocr), long-context vision-language understanding.")             
+            gr.Markdown("> [coreOCR-7B-050325-preview](https://huggingface.co/prithivMLmods/coreOCR-7B-050325-preview): model is a fine-tuned version of qwen/qwen2-vl-7b, optimized for document-level optical character recognition (ocr), long-context vision-language understanding.")  
+            gr.Markdown("> [SpaceOm](https://huggingface.co/remyxai/SpaceOm): SpaceOm, the reasoning traces in the spacethinker dataset average ~200 "thinking" tokens, so now included longer reasoning traces in the training data to help the model use more tokens in reasoning.")             
 
     image_submit.click(
         fn=generate_image,