Update app.py

app.py

@@ -13,7 +13,7 @@ from transformers import (
     AutoModelForCausalLM,
     AutoProcessor,
     TextIteratorStreamer,
-    AutoTokenizer, # Added for DeepSeek, though AutoProcessor is used
+    AutoTokenizer,
 )
 
 from gradio.themes import Soft
@@ -135,27 +135,24 @@ if os.path.exists(modeling_file_path):
     with open(modeling_file_path, 'r', encoding='utf-8') as f:
         input_code = f.read()
 
-    # The problematic import line
     original_import = "from transformers.models.llama.modeling_llama import (\n    LlamaAttention,\n    LlamaFlashAttention2\n)"
-    
     if original_import in input_code:
-        # Replace with a safe version that handles the ImportError
         safe_import = """from transformers.models.llama.modeling_llama import (
     LlamaAttention
 )
 try:
     from transformers.models.llama.modeling_llama import LlamaFlashAttention2
 except ImportError:
-    print("Warning: `LlamaFlashAttention2` not found. Falling back to `LlamaAttention`.")
     LlamaFlashAttention2 = LlamaAttention"""
-        
         patched_code = input_code.replace(original_import, safe_import)
-        
         with open(modeling_file_path, 'w', encoding='utf-8') as f:
             f.write(patched_code)
         print("Patched modeling_deepseekv2.py successfully.")
 sys.path.append(model_path_s_local)
 
+# --- NEW: Import the specific model class for DeepSeek-OCR ---
+from modeling_deepseekocr import DeepseekOCRForCausalLM
+
 
 # --- Model Loading ---
 
@@ -186,12 +183,13 @@ model_d = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 ).eval()
 
-# Load DeepSeek-OCR from the local, patched directory
+# Load DeepSeek-OCR from the local, patched directory using its specific class
 MODEL_PATH_S = model_path_s_local
 processor_s = AutoProcessor.from_pretrained(MODEL_PATH_S, trust_remote_code=True)
-model_s = AutoModelForCausalLM.from_pretrained(
+# --- MODIFIED: Use the specific class instead of AutoModelForCausalLM ---
+model_s = DeepseekOCRForCausalLM.from_pretrained(
     MODEL_PATH_S,
-    _attn_implementation='flash_attention_2',
+    _attn_implementation='eager',
     torch_dtype=torch.bfloat16,
     trust_remote_code=True,
 ).to(device).eval()
@@ -221,11 +219,7 @@ def generate_image(model_name: str, text: str, image: Image.Image,
 
     images = [image.convert("RGB")]
 
-    # For DeepSeek-OCR, the recommended prompt format is slightly different
     if model_name == "DeepSeek-OCR":
-        # Using a format found in documentation for better performance
-        # Note: The processor is expected to handle the full templating.
-        # This approach follows the user's implementation.
         messages = [
             {"role": "user", "content": f"<image>\n<|grounding|>{text}"}
         ]