Update app.py

app.py

@@ -22,11 +22,6 @@ from transformers import (
 )
 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
@@ -61,16 +56,6 @@ 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.
@@ -111,9 +96,6 @@ 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
@@ -168,9 +150,6 @@ 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
@@ -269,7 +248,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", "SmolVLM2"],
+                choices=["VIREX-062225-exp", "DREX-062225-exp", "olmOCR-7B-0225"],
                 label="Select Model",
                 value="VIREX-062225-exp"
             )