Add error handling and fallback device mapping for model loading

app.py

@@ -72,23 +72,54 @@ class Model:
     def load_models(self):
         import os
         os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        os.environ["TRANSFORMERS_OFFLINE"] = "0"
         
-        self.colqwen2_model = ColQwen2.from_pretrained(
-            "vidore/colqwen2-v0.1",
-            torch_dtype=torch.bfloat16,
-            device_map="cuda:0",
-            trust_remote_code=True,
-        )
+        # Load ColQwen2 with explicit configuration
+        try:
+            self.colqwen2_model = ColQwen2.from_pretrained(
+                "vidore/colqwen2-v0.1",
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            )
+        except Exception as e:
+            print(f"Error loading ColQwen2: {e}")
+            # Fallback to CPU loading then move to GPU
+            self.colqwen2_model = ColQwen2.from_pretrained(
+                "vidore/colqwen2-v0.1",
+                torch_dtype=torch.bfloat16,
+                device_map=None,
+                trust_remote_code=True,
+            )
+            self.colqwen2_model = self.colqwen2_model.to("cuda:0")
+            
         self.colqwen2_processor = ColQwen2Processor.from_pretrained(
             "vidore/colqwen2-v0.1"
         )
-        self.qwen2_vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
-            MODEL_NAME,
-            revision=MODEL_REVISION,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-        )
-        self.qwen2_vl_model.to("cuda:0")
+        
+        # Load Qwen2-VL with explicit configuration  
+        try:
+            self.qwen2_vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
+                MODEL_NAME,
+                revision=MODEL_REVISION,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=True,
+                device_map="auto",
+                low_cpu_mem_usage=True,
+            )
+        except Exception as e:
+            print(f"Error loading Qwen2VL: {e}")
+            # Fallback approach
+            self.qwen2_vl_model = Qwen2VLForConditionalGeneration.from_pretrained(
+                MODEL_NAME,
+                revision=MODEL_REVISION,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=True,
+                device_map=None,
+            )
+            self.qwen2_vl_model = self.qwen2_vl_model.to("cuda:0")
+            
         self.qwen2_vl_processor = AutoProcessor.from_pretrained(
             MODEL_NAME, 
             revision=MODEL_REVISION,