Improve model initialization with CPU-first approach

app.py

@@ -24,10 +24,6 @@ MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
 CACHE_DIR = "/home/user/.cache/huggingface"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Enable ZeroGPU features
-os.environ["ZERO_GPU"] = "1"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
 class Review:
     def __init__(self, code: str, language: str, suggestions: str):
         self.code = code
@@ -64,28 +60,36 @@ class CodeReviewer:
             )
             
             logger.info("Loading model...")
-            # Initialize model with ZeroGPU configuration
+            # Initialize model with specific configuration
             model_kwargs = {
-                "torch_dtype": torch.float16,
                 "trust_remote_code": True,
                 "low_cpu_mem_usage": True,
                 "cache_dir": CACHE_DIR,
-                "token": HF_TOKEN,
-                "device_map": {"": 0}  # Use first GPU
+                "token": HF_TOKEN
             }
             
-            # Load model with ZeroGPU support
-            self.model = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME,
-                **model_kwargs
-            )
+            # First try loading with CPU
+            try:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    MODEL_NAME,
+                    device_map=None,  # Start with no device map
+                    **model_kwargs
+                )
+                self.device = torch.device("cpu")
+                logger.info("Model loaded on CPU, will attempt GPU transfer")
+            except Exception as e1:
+                logger.error(f"Failed to load model on CPU: {e1}")
+                raise
             
-            # Set device based on ZeroGPU availability
-            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-            if self.device.type == "cpu":
-                logger.warning("Running on CPU - ZeroGPU not initialized")
-            else:
-                logger.info("ZeroGPU initialized successfully")
+            # Try moving to GPU if available
+            try:
+                if torch.cuda.is_available():
+                    logger.info("Moving model to GPU")
+                    self.model = self.model.to("cuda")
+                    self.device = torch.device("cuda")
+            except Exception as e2:
+                logger.warning(f"Could not move model to GPU: {e2}")
+                logger.info("Continuing with CPU")
             
             logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e: