Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 from PIL import Image
 from sd_parsers import ParserManager
 from torchvision import transforms
-from transformers import CLIPProcessor, CLIPModel, Blip2Processor, Blip2ForConditionalGeneration
+from transformers import CLIPProcessor, CLIPModel, Blip2Processor, Blip2ForConditionalGeneration, BitsAndBytesConfig
 import lpips
 import piq
 import plotly.express as px
@@ -18,31 +18,33 @@ import plotly.express as px
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # CLIP for prompt alignment & aesthetics
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+clip_model = CLIPModel.from_pretrained(
+    "openai/clip-vit-base-patch32"
+).to(device)
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# BLIP-2 for caption generation (processor without .to)
+# BLIP-2 for caption generation (8-bit quantized / fp8 proxy)
+bnb_config = BitsAndBytesConfig(load_in_8bit=True)
 blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
 blip_model = Blip2ForConditionalGeneration.from_pretrained(
-    "Salesforce/blip2-flan-t5-xl", torch_dtype=torch.float16
-).to(device)
+    "Salesforce/blip2-flan-t5-xl",
+    quantization_config=bnb_config,
+    device_map="auto"
+)
 
-# LPIPS for diversity\ nlpips_model = lpips.LPIPS(net='alex').to(device)
+# LPIPS for diversity
+lpips_model = lpips.LPIPS(net='alex').to(device)
 
 # --------------------
 # Helper Functions
 # --------------------
 
-def extract_metadata(image_bytes):
-    """Extracts prompt and model name from image bytes using sd-parsers."""
+def extract_metadata(file):
+    """Extract prompt and model name using sd-parsers from file path."""
     parser = ParserManager()
-    tmp_path = "temp.png"
-    with open(tmp_path, 'wb') as tmp:
-        tmp.write(image_bytes)
-    info = parser.parse(tmp_path)
+    info = parser.parse(file.name)
     prompt = info.prompts[0].value if info.prompts else ''
     model_name = info.model_name or ''
-    os.remove(tmp_path)
     return prompt, model_name
 
 # Image preprocessing transform
@@ -90,9 +92,9 @@ def analyze_images(files):
     imgs_by_model = {}
 
     for f in files:
-        image_bytes = f.read()
-        img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
-        prompt, model = extract_metadata(image_bytes)
+        # use f.name path instead of read() to avoid NamedString issues
+        img = Image.open(f.name).convert('RGB')
+        prompt, model = extract_metadata(f)
 
         clip_score = compute_clip_score(img, prompt)
         cap_sim = compute_caption_similarity(img, prompt)