Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 OCR Application with Multiple Models including DeepSeek OCR
-Fixed version with @spaces.GPU decorator for Hugging Face Spaces
+Fixed version with proper chat template handling for DeepSeek OCR
 """
 
 import os
@@ -97,7 +97,7 @@ except Exception as e:
     processor_m = None
     print(f"✗ olmOCR-2-7B-1025: Failed to load - {str(e)}")
 
-# Load DeepSeek-OCR
+# Load DeepSeek-OCR with custom chat template
 try:
     MODEL_ID_DS = "deepseek-ai/deepseek-ocr"
     processor_ds = AutoProcessor.from_pretrained(MODEL_ID_DS, trust_remote_code=True)
@@ -106,7 +106,14 @@ try:
         trust_remote_code=True,
         torch_dtype=torch.float16
     ).eval()
-    print("✓ DeepSeek-OCR loaded")
+    
+    # Set a default chat template for DeepSeek OCR if it doesn't have one
+    if not hasattr(processor_ds.tokenizer, 'chat_template') or processor_ds.tokenizer.chat_template is None:
+        # Use a standard Qwen-style chat template
+        processor_ds.tokenizer.chat_template = "{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}<|im_end|>\n{% endfor %}<|im_start|>assistant\n"
+        print("✓ DeepSeek-OCR loaded (with custom chat template)")
+    else:
+        print("✓ DeepSeek-OCR loaded")
 except Exception as e:
     model_ds = None
     processor_ds = None
@@ -189,12 +196,17 @@ def generate_image(model_name: str, text: str, image: Image.Image,
             ]
         }]
         
-        # Apply chat template
-        prompt_full = processor.apply_chat_template(
-            messages, 
-            tokenize=False, 
-            add_generation_prompt=True
-        )
+        # Apply chat template with error handling
+        try:
+            prompt_full = processor.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
+        except Exception as template_error:
+            # Fallback: create a simple prompt without chat template
+            print(f"Chat template error: {template_error}. Using fallback prompt.")
+            prompt_full = f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
 
         # Process inputs
         inputs = processor(
@@ -239,6 +251,9 @@ def generate_image(model_name: str, text: str, image: Image.Image,
         
     except Exception as e:
         error_msg = f"Error during generation: {str(e)}"
+        print(f"Full error: {e}")
+        import traceback
+        traceback.print_exc()
         yield error_msg, error_msg