Update app.py

app.py

@@ -118,22 +118,22 @@ MODEL_ID_3B = "nanonets/Nanonets-OCR2-3B"
 processor_3b = AutoProcessor.from_pretrained(MODEL_ID_3B, trust_remote_code=True)
 model_3b = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID_3B,
-    torch_dtype="auto",
-    device_map="auto",
+    dtype=torch.float16,
+    #device_map="auto",
     trust_remote_code=True,
     attn_implementation="flash_attention_2"
-).eval()
+).to(device).eval()
 
 # Load Nanonets-OCR2-1.5B-exp
 MODEL_ID_1_5B = "nanonets/Nanonets-OCR2-1.5B-exp"
 processor_1_5b = AutoProcessor.from_pretrained(MODEL_ID_1_5B, trust_remote_code=True)
 model_1_5b = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID_1_5B,
-    torch_dtype="auto",
-    device_map="auto",
+    dtype=torch.float16,
+    #device_map="auto",
     trust_remote_code=True,
     attn_implementation="flash_attention_2"
-).eval()
+).to(device).eval()
 
 @spaces.GPU
 def generate_image(model_name: str, text: str, image: Image.Image,