Update app.py

app.py

@@ -23,14 +23,22 @@ clip_model = CLIPModel.from_pretrained(
 ).to(device)
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# BLIP-2 for caption generation (8-bit quantized / fp8 proxy)
-bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+# BLIP-2 for caption generation: 8-bit if GPU available, else float16
 blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
-blip_model = Blip2ForConditionalGeneration.from_pretrained(
-    "Salesforce/blip2-flan-t5-xl",
-    quantization_config=bnb_config,
-    device_map="auto"
-)
+if torch.cuda.is_available():
+    bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+    blip_model = Blip2ForConditionalGeneration.from_pretrained(
+        "Salesforce/blip2-flan-t5-xl",
+        quantization_config=bnb_config,
+        device_map="auto"
+    )
+else:
+    # CPU-only environment: load half precision
+    blip_model = Blip2ForConditionalGeneration.from_pretrained(
+        "Salesforce/blip2-flan-t5-xl",
+        torch_dtype=torch.float16
+    )
+    blip_model.to(device)
 
 # LPIPS for diversity
 lpips_model = lpips.LPIPS(net='alex').to(device)
@@ -92,7 +100,6 @@ def analyze_images(files):
     imgs_by_model = {}
 
     for f in files:
-        # use f.name path instead of read() to avoid NamedString issues
         img = Image.open(f.name).convert('RGB')
         prompt, model = extract_metadata(f)
 
@@ -114,7 +121,6 @@ def analyze_images(files):
 
     df = pd.DataFrame(records)
 
-    # Diversity per model
     diversity = {}
     for model, imgs in imgs_by_model.items():
         if len(imgs) < 2: