Update app.py

app.py

@@ -24,6 +24,9 @@ from transformers import (
 
 from transformers.image_utils import load_image
 
+import subprocess
+subprocess.run('pip install flash-attn==2.6.3 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
 # Constants for text generation
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -58,13 +61,15 @@ model_k = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load Imgscope-OCR-2B-0527
-MODEL_ID_Y = "prithivMLmods/Imgscope-OCR-2B-0527"
-processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
-model_y = Qwen2VLForConditionalGeneration.from_pretrained(
+# 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(
     MODEL_ID_Y,
     trust_remote_code=True,
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",
+    revision='1f0fdea7f5b19532a750be109b19072d719b8177'
 ).to(device).eval()
 
 
@@ -108,8 +113,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 == "Imgscope-OCR-2B-0527":
-        processor = processor_y
+    elif model_name == "llama-nemoretriever-colembed-1b-v1":
+        #processor = processor_y
         model = model_y
     else:
         yield "Invalid model selected."
@@ -165,8 +170,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 == "Imgscope-OCR-2B-0527":
-        processor = processor_y
+    elif model_name == "llama-nemoretriever-colembed-1b-v1":
+        #processor = processor_y
         model = model_y
     else:
         yield "Invalid model selected."
@@ -269,7 +274,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", "Imgscope-OCR-2B-0527"],
+                choices=["SkyCaptioner-V1", "SpaceThinker-3B", "coreOCR-7B-050325-preview", "llama-nemoretriever-colembed-1b-v1"],
                 label="Select Model",
                 value="SkyCaptioner-V1"
             )
@@ -277,7 +282,6 @@ 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("> [Imgscope-OCR-2B-0527](https://huggingface.co/prithivMLmods/Imgscope-OCR-2B-0527): fine-tuned version of qwen2-vl-2b-instruct, specifically optimized for messy handwriting recognition, document ocr, realistic handwritten ocr, and math problem solving with latex formatting.")
             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.")             
 
     image_submit.click(