update

app_high_res.py

@@ -11,9 +11,12 @@ import functools
 from transformers import AutoProcessor, Idefics2ForConditionalGeneration
 from models.conversation import conv_templates
 from typing import List
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
 processor = AutoProcessor.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-high-res-35k-mantis-2epoch_4096")
-model = Idefics2ForConditionalGeneration.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-high-res-35k-mantis-2epoch_4096", torch_dtype=torch.bfloat16)
-# model = model.to("cuda")
+model = Idefics2ForConditionalGeneration.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-high-res-35k-mantis-2epoch_4096", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
 MAX_NUM_FRAMES = 24
 conv_template = conv_templates["idefics_2"]
 
@@ -61,8 +64,7 @@ all the frames of video are as follows:
 """
 @spaces.GPU(duration=60)
 def generate(text:str, images:List[Image.Image], history: List[dict], **kwargs):
-    global processor, model
-    model = model.to("cuda")
+    model.to("cuda")
     if not images:
         images = None