Update app.py

app.py

@@ -39,7 +39,7 @@ class ChatCompletionRequest(BaseModel):
     messages: List[Message]
 
 # ------------------------------------------------------------
-# 3.  Lazy MedGemma loader (CPU; float32)
+# 3.  Lazy MedGemma loader with memory optimization
 # ------------------------------------------------------------
 DEVICE = "cpu"
 DTYPE  = torch.float32
@@ -49,20 +49,23 @@ def get_medgemma():
     global medgemma_pipe
     if medgemma_pipe is None:
         try:
-            print("Loading MedGemma-4B-IT …")
+            print("🚀 Loading MedGemma-4B-IT with memory optimization...")
             print(f"Using cache directory: {CACHE_DIR}")
             
             medgemma_pipe = pipeline(
                 "text-generation",
                 model="google/medgemma-4b-it",
-                torch_dtype=DTYPE,
-                device_map="cpu",  # Force CPU to avoid device issues
+                torch_dtype=torch.float16,  # Use float16 to reduce memory
+                device_map="auto",
                 token=HF_TOKEN,
                 cache_dir=CACHE_DIR,
                 trust_remote_code=True,
-                local_files_only=False,  # Allow downloads
+                # Memory optimization
+                low_cpu_mem_usage=True,
+                max_memory={0: "6GB", "cpu": "8GB"},  # Limit memory usage
             )
-            print("✅ MedGemma loaded successfully")
+            print("✅ MedGemma loaded successfully!")
+            
         except Exception as e:
             print(f"❌ Error loading MedGemma: {e}")
             print(f"Cache directory exists: {os.path.exists(CACHE_DIR)}")