Upload 2 files

app.py

@@ -54,7 +54,7 @@ def encode_video(video_path, max_num_frames=32):
         return []
 
 def load_model():
-    """Load MiniCPM-o 2.6 model and tokenizer"""
+    """Load MiniCPM-o 2.6 model and tokenizer with fallback attention implementation"""
     global model, tokenizer
     
     if model is not None and tokenizer is not None:
@@ -69,11 +69,20 @@ def load_model():
         
         print(f"Loading on device: {device} with dtype: {torch_dtype}")
         
+        # Try to determine best attention implementation
+        attn_implementation = 'sdpa'  # Default fallback
+        try:
+            import flash_attn
+            attn_implementation = 'flash_attention_2'
+            print("✅ Using Flash Attention 2 for optimal performance")
+        except ImportError:
+            print("⚠️ Flash Attention not available, using SDPA (still good performance)")
+        
         # Load model with memory optimization for Spaces
         model = AutoModel.from_pretrained(
             model_path,
             trust_remote_code=True,
-            attn_implementation='sdpa',
+            attn_implementation=attn_implementation,
             torch_dtype=torch_dtype,
             device_map="auto",
             offload_buffers=True,  # Enable buffer offloading for GPU memory
@@ -89,8 +98,34 @@ def load_model():
         return model, tokenizer
         
     except Exception as e:
-        print(f"❌ Error loading model: {e}")
-        raise e
+        # If flash_attn fails, try with SDPA as fallback
+        if "flash_attn" in str(e):
+            print("⚠️ Flash Attention failed, retrying with SDPA...")
+            try:
+                model = AutoModel.from_pretrained(
+                    model_path,
+                    trust_remote_code=True,
+                    attn_implementation='sdpa',
+                    torch_dtype=torch_dtype,
+                    device_map="auto",
+                    offload_buffers=True,
+                    low_cpu_mem_usage=True
+                )
+                
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path, 
+                    trust_remote_code=True
+                )
+                
+                print("✅ Model loaded successfully with SDPA fallback!")
+                return model, tokenizer
+                
+            except Exception as fallback_e:
+                print(f"❌ Error loading model with fallback: {fallback_e}")
+                raise fallback_e
+        else:
+            print(f"❌ Error loading model: {e}")
+            raise e
 
 def analyze_video(video_file, prompt, max_frames):
     """Analyze video using MiniCPM-o 2.6"""

requirements.txt

@@ -18,4 +18,5 @@ numpy>=1.24.0
 huggingface_hub>=0.33.2
 protobuf==3.20.3
 einops>=0.8.0
-einx>=0.3.0 
+einx>=0.3.0
+flash-attn==2.8.0.post2