Update app.py

app.py

@@ -5,13 +5,20 @@ from transformers import (
     AutoModelForImageTextToText,
     AutoTokenizer, 
     AutoProcessor,
-    BitsAndBytesConfig,
     pipeline
 )
 from PIL import Image
 import os
 import spaces
 
+# Try to import bitsandbytes for quantization (optional)
+try:
+    from transformers import BitsAndBytesConfig
+    QUANTIZATION_AVAILABLE = True
+except ImportError:
+    QUANTIZATION_AVAILABLE = False
+    print("⚠️ bitsandbytes not available. Quantization will be disabled.")
+
 # Configuration
 MODEL_4B = "google/medgemma-4b-it"
 MODEL_27B = "google/medgemma-27b-text-it"
@@ -31,8 +38,11 @@ class MedGemmaApp:
             "device_map": "auto",
         }
         
-        if use_quantization:
+        # Only add quantization if available and requested
+        if use_quantization and QUANTIZATION_AVAILABLE:
             model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True)
+        elif use_quantization and not QUANTIZATION_AVAILABLE:
+            print("⚠️ Quantization requested but bitsandbytes not available. Loading without quantization.")
         
         return model_kwargs
     
@@ -172,9 +182,10 @@ with gr.Blocks(title="MedGemma Medical AI Assistant", theme=gr.themes.Soft()) as
             )
             
             use_quantization = gr.Checkbox(
-                value=True,
-                label="Use 4-bit Quantization",
-                info="Reduces memory usage (recommended)"
+                value=QUANTIZATION_AVAILABLE,
+                label="Use 4-bit Quantization" + ("" if QUANTIZATION_AVAILABLE else " (Unavailable)"),
+                info="Reduces memory usage" + ("" if QUANTIZATION_AVAILABLE else " - bitsandbytes not installed"),
+                interactive=QUANTIZATION_AVAILABLE
             )
             
             load_btn = gr.Button("🚀 Load Model", variant="primary")