Update app.py

app.py

@@ -16,11 +16,17 @@ import cv2
 from transformers import (
     Qwen2VLForConditionalGeneration,
     Qwen2_5_VLForConditionalGeneration,
+    AutoModelForImageTextToText,
     AutoProcessor,
     TextIteratorStreamer,
 )
 from transformers.image_utils import load_image
 
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+from io import BytesIO
+
 # Constants for text generation
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -55,6 +61,16 @@ model_o = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
+# Load SmolVLM2-2.2B-Instruct
+MODEL_ID_W = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
+processor_w = AutoProcessor.from_pretrained(MODEL_ID_W, trust_remote_code=True)
+model_w= AutoModelForImageTextToText.from_pretrained(
+    MODEL_ID_W,
+    trust_remote_code=True,
+    _attn_implementation="flash_attention_2",
+    torch_dtype=torch.float16
+).to(device).eval()
+
 def downsample_video(video_path):
     """
     Downsamples the video to evenly spaced frames.
@@ -95,6 +111,9 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "olmOCR-7B-0225":
         processor = processor_o
         model = model_o
+    elif model_name == "SmolVLM2":
+        processor = processor_w
+        model = model_w
     else:
         yield "Invalid model selected."
         return
@@ -149,6 +168,9 @@ def generate_video(model_name: str, text: str, video_path: str,
     elif model_name == "olmOCR-7B-0225":
         processor = processor_o
         model = model_o
+    elif model_name == "SmolVLM2":
+        processor = processor_w
+        model = model_w
     else:
         yield "Invalid model selected."
         return
@@ -247,7 +269,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=["VIREX-062225-exp", "DREX-062225-exp", "olmOCR-7B-0225"],
+                choices=["VIREX-062225-exp", "DREX-062225-exp", "olmOCR-7B-0225", "SmolVLM2"],
                 label="Select Model",
                 value="VIREX-062225-exp"
             )