Configure app for ZeroGPU dynamic GPU allocation

app.py

@@ -24,6 +24,10 @@ 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
@@ -36,7 +40,7 @@ class CodeReviewer:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        self.device = "cpu"  # Default to CPU
+        self.device = None
         self.review_history: List[Review] = []
         self.metrics = {
             'total_reviews': 0,
@@ -60,21 +64,28 @@ class CodeReviewer:
             )
             
             logger.info("Loading model...")
-            # Initialize model with specific configuration
+            # Initialize model with ZeroGPU configuration
             model_kwargs = {
-                "torch_dtype": torch.float32,  # Use float32 for CPU
+                "torch_dtype": torch.float16,
                 "trust_remote_code": True,
                 "low_cpu_mem_usage": True,
                 "cache_dir": CACHE_DIR,
-                "token": HF_TOKEN
+                "token": HF_TOKEN,
+                "device_map": {"": 0}  # Use first GPU
             }
             
-            # Load model directly to CPU
+            # Load model with ZeroGPU support
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_NAME,
-                device_map=None,  # Don't use device_map
                 **model_kwargs
-            ).to(self.device)
+            )
+            
+            # 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")
             
             logger.info(f"Model loaded successfully on {self.device}")
         except Exception as e:
@@ -146,6 +157,11 @@ Code:
             # Update metrics
             self.update_metrics(review)
             
+            # Clear GPU memory if using CUDA
+            if self.device.type == "cuda":
+                del inputs, outputs
+                torch.cuda.empty_cache()
+            
             return suggestions
             
         except Exception as e: