Update app.py

app.py

@@ -5,31 +5,41 @@ import logging
 from typing import Optional, Union
 import os
 import spaces
-from dotenv import load_dotenv
 
-load_dotenv()
+# Disable torch compilation to avoid dynamo issues
+torch._dynamo.config.disable = True
+torch.backends.cudnn.allow_tf32 = True
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class AtlasOCR:
-    def __init__(self, model_name: str = "atlasia/AtlasOCR", max_tokens: int = 2000):
+    def __init__(self, model_name: str = "atlasia/AtlasOCR-v0", max_tokens: int = 2000):
         """Initialize the AtlasOCR model with proper error handling."""
         try:
             from unsloth import FastVisionModel
             
             logger.info(f"Loading model: {model_name}")
-            self.model, self.processor = FastVisionModel.from_pretrained(
-                model_name,
-                device_map="cuda",
-                load_in_4bit=True,
-                use_gradient_checkpointing="unsloth",
-                token=os.environ["HF_API_KEY"]
-            )
+            
+            # Disable compilation for the model
+            with torch._dynamo.config.patch(disable=True):
+                self.model, self.processor = FastVisionModel.from_pretrained(
+                    model_name,
+                    device_map="auto",
+                    load_in_4bit=True,
+                    use_gradient_checkpointing="unsloth",
+                    torch_dtype=torch.float16
+                )
+            
+            # Ensure model is not compiled
+            if hasattr(self.model, '_dynamo_compile'):
+                self.model._dynamo_compile = False
+                
             self.max_tokens = max_tokens
             self.prompt = ""
-            logger.info("Model loaded successfully")
+            self.device = next(self.model.parameters()).device
+            logger.info(f"Model loaded successfully on device: {self.device}")
             
         except ImportError:
             logger.error("unsloth not found. Please install it: pip install unsloth")
@@ -81,22 +91,30 @@ class AtlasOCR:
             
             inputs = self.prepare_inputs(image)
             
-            # Move inputs to GPU if available
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            inputs = {k: v.to(device) if hasattr(v, 'to') else v for k, v in inputs.items()}
+            # Move inputs to the same device as model with explicit device handling
+            device = self.device
+            logger.info(f"Moving inputs to device: {device}")
+            
+            # Manually move each tensor to device
+            for key in inputs:
+                if hasattr(inputs[key], 'to'):
+                    inputs[key] = inputs[key].to(device)
             
-            # Ensure attention_mask is float32
+            # Ensure attention_mask is float32 and on correct device
             if 'attention_mask' in inputs:
-                inputs['attention_mask'] = inputs['attention_mask'].to(torch.float32)
+                inputs['attention_mask'] = inputs['attention_mask'].to(dtype=torch.float32, device=device)
             
             logger.info(f"Generating text with max_tokens={self.max_tokens}")
-            with torch.no_grad():
+            
+            # Disable compilation during generation
+            with torch.no_grad(), torch._dynamo.config.patch(disable=True):
                 generated_ids = self.model.generate(
                     **inputs, 
                     max_new_tokens=self.max_tokens, 
                     use_cache=True,
                     do_sample=False,
-                    temperature=0.1
+                    temperature=0.1,
+                    pad_token_id=self.processor.tokenizer.eos_token_id
                 )
             
             generated_ids_trimmed = [