Update app.py

app.py

@@ -20,7 +20,6 @@ def free_gpu_cache():
 # =========================
 # MODELS
 # =========================
-# Image generation
 gen_pipe = DiffusionPipeline.from_pretrained(
     "stabilityai/sdxl-turbo",
     torch_dtype=torch.float16 if device=="cuda" else torch.float32
@@ -31,7 +30,6 @@ dreamshaper_pipe = DiffusionPipeline.from_pretrained(
     torch_dtype=torch.float16 if device=="cuda" else torch.float32
 ).to(device)
 
-# Captioning
 captioner = pipeline(
     "image-to-text",
     model="Salesforce/blip-image-captioning-large",
@@ -39,7 +37,6 @@ captioner = pipeline(
     generate_kwargs={"max_new_tokens":256, "num_beams":5, "temperature":0.7}
 )
 
-# NLP
 sentiment_model = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english",
                            device=0 if device=="cuda" else -1)
 ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english",
@@ -47,16 +44,13 @@ ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-engl
 topic_model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli",
                        device=0 if device=="cuda" else -1)
 
-# VQA
 vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
 vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to("cpu")
 
-# Metrics
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 lpips_model = lpips.LPIPS(net='alex').to(device)
 lpips_transform = T.Compose([T.ToTensor(), T.Resize((256,256))])
 
-# Styles
 style_map = {
     "Photorealistic": "photorealistic, ultra-detailed, 8k, cinematic lighting",
     "Real Life": "natural lighting, true-to-life colors, DSLR",
@@ -74,9 +68,7 @@ style_map = {
 # IMAGE GENERATION FUNCTIONS
 # =========================
 def generate_image_with_enhancer(base_caption, enhancer, negative, seed, style, images):
-    images = images or []
-    base_caption = base_caption or ""
-    enhancer = enhancer or ""
+    images = images or [None, None, None]
     final_prompt = f"{base_caption}, {enhancer}".strip(", ")
     final_prompt = f"{final_prompt}, {style_map.get(style,'')}".strip(", ")
     try:
@@ -92,14 +84,12 @@ def generate_image_with_enhancer(base_caption, enhancer, negative, seed, style,
         print("SD Turbo failed:", e)
         img = None
     if img:
-        images.append(img)
+        images[1] = img  # Always put SD-Turbo at index 1
     free_gpu_cache()
     return img, images
 
 def generate_dreamshaper_with_enhancer(base_caption, enhancer, negative, seed, style, images):
-    images = images or []
-    base_caption = base_caption or ""
-    enhancer = enhancer or ""
+    images = images or [None, None, None]
     final_prompt = f"{base_caption}, {enhancer}".strip(", ")
     final_prompt = f"{final_prompt}, {style_map.get(style,'')}".strip(", ")
     try:
@@ -115,7 +105,7 @@ def generate_dreamshaper_with_enhancer(base_caption, enhancer, negative, seed, s
         print("DreamShaper failed:", e)
         img = None
     if img:
-        images.append(img)
+        images[2] = img  # Always put DreamShaper at index 2
     free_gpu_cache()
     return img, images
 
@@ -133,7 +123,7 @@ def caption_for_image(img):
 # VQA
 # =========================
 def answer_vqa(question, image):
-    if not image or not question.strip():
+    if image is None or not question.strip():
         return "Provide image + question."
     try:
         inputs_raw = vqa_processor(images=image, text=question, return_tensors="pt")
@@ -175,14 +165,13 @@ def compute_metrics(images, captions, i1, i2):
     else:
         bert_f1 = 0.0
 
-    return clip_sim, lp, bert_f1
+    return f"CLIP: {clip_sim:.2f}\nLPIPS: {lp:.2f}\nBERTScore F1: {bert_f1:.2f}"
 
 # =========================
 # GRADIO UI BUILD
 # =========================
 def build_full_ui():
     with gr.Blocks(title="Multimodal AI Image Studio") as demo:
-        # --- CSS Styling ---
         gr.HTML("""
         <style>
         .heading-orange h2, .heading-orange h3 { color: #ff5500 !important; }
@@ -196,13 +185,10 @@ def build_full_ui():
         </style>
         """)
 
-        # --- States ---
         images_state = gr.State([None, None, None])
         captions_state = gr.State(["", "", ""])
 
-        # =========================
-        # Section 1: Upload Reference Image
-        # =========================
+        # --- Upload Section ---
         gr.Markdown("## 1️⃣ Upload Reference Image", elem_classes="heading-orange")
         with gr.Row(elem_classes="equal-height-row"):
             with gr.Column(scale=1):
@@ -213,24 +199,17 @@ def build_full_ui():
                 upload_preview = gr.Image(label="Uploaded Image", interactive=False)
                 caption_out = gr.Markdown(label="Generated Caption")
 
-        # Upload & caption function
         def upload_and_caption(img, images_state, captions_state):
             if img is None:
                 return None, "No image uploaded.", images_state, captions_state
             images_state[0] = img
-            try:
-                cap = caption_for_image(img)
-            except:
-                cap = "Caption failed."
-            captions_state[0] = cap
-            return img, cap, images_state, captions_state
+            captions_state[0] = caption_for_image(img)
+            return img, captions_state[0], images_state, captions_state
 
         upload_btn.click(upload_and_caption, inputs=[upload_input, images_state, captions_state],
                          outputs=[upload_preview, caption_out, images_state, captions_state])
 
-        # =========================
-        # Section 2: Generate SD-Turbo & DreamShaper
-        # =========================
+        # --- Generate SD-Turbo & DreamShaper ---
         gr.Markdown("## 2️⃣ Generate Images from Caption", elem_classes="heading-orange")
         with gr.Row():
             with gr.Column(scale=1):
@@ -240,16 +219,14 @@ def build_full_ui():
                 ds_btn = gr.Button("Generate DreamShaper Image", elem_classes="orange-btn")
                 ds_preview = gr.Image(label="DreamShaper Image", interactive=False)
 
-        # Generate SD-Turbo
         def generate_sd(caption, enhancer, images_state, captions_state):
-            img, images_state = generate_image_with_enhancer(caption, enhancer, negative="", seed=42, style="Photorealistic", images=images_state)
+            img, images_state = generate_image_with_enhancer(caption, enhancer, "", 42, "Photorealistic", images_state)
             if img:
                 captions_state[1] = caption_for_image(img)
             return img, images_state, captions_state
 
-        # Generate DreamShaper
         def generate_ds(caption, enhancer, images_state, captions_state):
-            img, images_state = generate_dreamshaper_with_enhancer(caption, enhancer, negative="", seed=123, style="Photorealistic", images=images_state)
+            img, images_state = generate_dreamshaper_with_enhancer(caption, enhancer, "", 123, "Photorealistic", images_state)
             if img:
                 captions_state[2] = caption_for_image(img)
             return img, images_state, captions_state
@@ -259,79 +236,29 @@ def build_full_ui():
         ds_btn.click(generate_ds, inputs=[caption_out, enhancer_box, images_state, captions_state],
                      outputs=[ds_preview, images_state, captions_state])
 
-        # =========================
-        # Section 3: Compute Pairwise Metrics (Side-by-Side)
-        # =========================
+        # --- Compute Metrics ---
         gr.Markdown("## 3️⃣ Compute Pairwise Metrics", elem_classes="heading-orange")
         metrics_btn = gr.Button("Compute Metrics for All Pairs", elem_classes="teal-btn")
         metrics_spinner = gr.HTML("<div style='height:4px;'></div>")
-        metrics_out = gr.HTML()
+        metrics_A = gr.Markdown()
+        metrics_B = gr.Markdown()
+        metrics_C = gr.Markdown()
 
         def compute_metrics_ui(images, captions):
-            yield "<div class='loading-line'></div>", ""
+            yield "<div class='loading-line'></div>", "", "", ""
             if any(i is None for i in images):
-                yield "All three images and captions are required."
+                msg = "All three images and captions are required."
+                yield "", msg, msg, msg
             else:
-                try:
-                    A = compute_metrics(images, captions, 0, 1)
-                    B = compute_metrics(images, captions, 0, 2)
-                    C = compute_metrics(images, captions, 1, 2)
-                    def fmt(m):
-                        return f"CLIP: {m[0]:.3f}<br>LPIPS: {m[1]:.3f}<br>BERTScore F1: {m[2]:.3f}"
-                    html = f"""
-                    <div style='display:flex; gap:40px; justify-content:space-around;'>
-                        <div style='text-align:center;'><b>Metrics A</b><br>{fmt(A)}</div>
-                        <div style='text-align:center;'><b>Metrics B</b><br>{fmt(B)}</div>
-                        <div style='text-align:center;'><b>Metrics C</b><br>{fmt(C)}</div>
-                    </div>
-                    """
-                    yield html
-                except Exception as e:
-                    print("Metrics error:", e)
-                    yield "Failed to compute metrics."
+                A = compute_metrics(images, captions, 0, 1)
+                B = compute_metrics(images, captions, 0, 2)
+                C = compute_metrics(images, captions, 1, 2)
+                yield "", f"**Reference ↔ SD-Turbo**\n{A}", f"**Reference ↔ DreamShaper**\n{B}", f"**SD-Turbo ↔ DreamShaper**\n{C}"
 
         metrics_btn.click(compute_metrics_ui, inputs=[images_state, captions_state],
-                          outputs=[metrics_out])
-
-        # =========================
-        # Section 4: NLP Analysis
-        # =========================
-        gr.Markdown("## 4️⃣ NLP Analysis of Captions", elem_classes="heading-orange")
-        nlp_btn = gr.Button("Analyze Captions", elem_classes="teal-btn")
-        nlp_spinner = gr.HTML("<div style='height:4px;'></div>")
-        nlp_out = gr.HTML()
-
-        def analyze_captions_ui(captions):
-            yield "<div class='loading-line'></div>", ""
-            if any(c=="" for c in captions):
-                yield "<b>All three captions are required for NLP analysis.</b>"
-            else:
-                labels = ["Reference", "SD-Turbo", "DreamShaper"]
-                blocks = []
-                for label, caption in zip(labels, captions):
-                    try:
-                        sentiment = "<br>".join([f"{s['label']}: {s['score']:.2f}" for s in sentiment_model(caption)])
-                    except:
-                        sentiment = "Sentiment failed."
-                    try:
-                        ents_list = ner_model(caption)
-                        ents = "<br>".join([f"{e.get('entity_group','')}: {e.get('word','')}" for e in ents_list]) or "None"
-                    except:
-                        ents = "NER failed."
-                    try:
-                        topics_data = topic_model(caption, candidate_labels=['people','animals','objects','food','nature'])
-                        topics = "<br>".join([f"{l}: {sc:.2f}" for l, sc in zip(topics_data.get('labels',[]), topics_data.get('scores',[]))])
-                    except:
-                        topics = "Topics failed."
-                    block = f"<div style='flex:1;padding:10px;min-width:250px;'><h3><u>{label}</u></h3><b>Sentiment</b><br>{sentiment}<br><br><b>Entities</b><br>{ents}<br><br><b>Topics</b><br>{topics}</div>"
-                    blocks.append(block)
-                yield f"<div style='display:flex; gap:20px; justify-content:space-between;'>{''.join(blocks)}</div>"
-
-        nlp_btn.click(analyze_captions_ui, inputs=[captions_state], outputs=[nlp_out])
+                          outputs=[metrics_spinner, metrics_A, metrics_B, metrics_C])
 
-        # =========================
-        # Section 5: Visual Question Answering
-        # =========================
+        # --- VQA ---
         gr.Markdown("## 5️⃣ Visual Question Answering (VQA)", elem_classes="heading-orange")
         with gr.Row():
             with gr.Column(scale=1):
@@ -341,19 +268,12 @@ def build_full_ui():
                 vqa_spinner = gr.HTML("<div style='height:4px;'></div>")
                 vqa_out = gr.Markdown(label="VQA Output")
 
-        def vqa_ui(question, image):
+        def vqa_ui(question, images_state):
             yield "<div class='loading-line'></div>", ""
-            if not question.strip() or image is None:
-                yield "Provide image + question."
-            else:
-                try:
-                    ans = answer_vqa(question, image)
-                    yield f"<b>Answer:</b> {ans}"
-                except Exception as e:
-                    print("VQA error:", e)
-                    yield "Could not determine the answer."
+            ans = answer_vqa(question, images_state[0])
+            yield "", ans
 
-        vqa_btn.click(vqa_ui, inputs=[vqa_input, upload_preview], outputs=[vqa_out])
+        vqa_btn.click(vqa_ui, inputs=[vqa_input, images_state], outputs=[vqa_spinner, vqa_out])
 
     return demo
 
@@ -361,269 +281,11 @@ def build_full_ui():
 demo = build_full_ui()
 demo.launch()
 
-
-
 """
-# Dump code
-# ==============================
-# Libraries
-# ==============================
-import torch
-import gradio as gr
-from PIL import Image
-from diffusers import DiffusionPipeline
-from transformers import pipeline, BlipProcessor, BlipForQuestionAnswering
-import lpips
-import clip
-from bert_score import score
-import torchvision.transforms as T
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def free_gpu_cache():
-    if device == "cuda":
-        torch.cuda.empty_cache()
-
-# ==============================
-# Load Models (HF-ready, memory safe)
-# ==============================
-# SDXL-Turbo
-gen_pipe = DiffusionPipeline.from_pretrained(
-    "stabilityai/sdxl-turbo",
-    torch_dtype=torch.float16 if device=="cuda" else torch.float32
-).to(device)
-
-# DreamShaper
-dreamshaper_pipe = DiffusionPipeline.from_pretrained(
-    "Lykon/dreamshaper-7",
-    torch_dtype=torch.float16 if device=="cuda" else torch.float32
-).to(device)
-
-# BLIP Captioning
-captioner = pipeline(
-    "image-to-text",
-    model="Salesforce/blip-image-captioning-large",
-    device=0 if device=="cuda" else -1,
-    generate_kwargs={"max_new_tokens":256, "num_beams":5, "temperature":0.7}
-)
-
-# Sentiment / NER / Topic
-sentiment_model = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english",
-                           device=0 if device=="cuda" else -1)
-ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english",
-                     aggregation_strategy="simple", device=0 if device=="cuda" else -1)
-topic_model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli",
-                       device=0 if device=="cuda" else -1)
-
-# BLIP VQA
-vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
-vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to("cpu")
-
-# CLIP / LPIPS
-clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
-lpips_model = lpips.LPIPS(net='alex').to(device)
-lpips_transform = T.Compose([T.ToTensor(), T.Resize((256,256))])
-
-# Style map
-style_map = {
-    "Photorealistic": "photorealistic, ultra-detailed, 8k, cinematic lighting",
-    "Real Life": "natural lighting, true-to-life colors, DSLR",
-    "Documentary": "documentary handheld muted colors",
-    "iPhone Camera": "iPhone photo natural HDR",
-    "Street Photography": "candid street ambient shadows",
-    "Cinematic": "cinematic lighting dramatic depth",
-    "Anime": "anime cel shaded vibrant",
-    "Watercolor": "watercolor soft wash art",
-    "Macro": "macro lens shallow DOF",
-    "Cyberpunk": "neon cyberpunk futuristic",
-}
-
-# ==============================
-# Functions
-# ==============================
-
-def generate_image(pipe, caption, enhancer, negative, seed, style):
-    final_prompt = f"{caption}, {enhancer}".strip(", ")
-    final_prompt = f"{final_prompt}, {style_map.get(style,'')}".strip(", ")
-
-    try:
-        seed = int(seed)
-    except:
-        seed = 42
-
-    generator = torch.Generator(device="cpu").manual_seed(seed)
-    img = None
-
-    try:
-        with torch.no_grad():
-            out = pipe(prompt=final_prompt, negative_prompt=negative, generator=generator, height=512, width=512)
-        img = out.images[0]
-    except Exception as e:
-        print(f"{pipe} generation failed:", e)
-
-    free_gpu_cache()
-    return img
-
-def caption_for_image(img):
-    try:
-        out = captioner(img)
-        return out[0]["generated_text"]
-    except:
-        return "Caption failed."
-
-def compute_metrics(images, captions, i1, i2):
-    img1, img2 = images[i1], images[i2]
-    cap1, cap2 = captions[i1], captions[i2]
-
-    # CLIP similarity
-    t1, t2 = clip_preprocess(img1).unsqueeze(0).to(device), clip_preprocess(img2).unsqueeze(0).to(device)
-    with torch.no_grad():
-        f1, f2 = clip_model.encode_image(t1), clip_model.encode_image(t2)
-        clip_sim = float(torch.cosine_similarity(f1, f2))
-
-    # LPIPS
-    L1 = (lpips_transform(img1).unsqueeze(0)*2 - 1).to(device)
-    L2 = (lpips_transform(img2).unsqueeze(0)*2 - 1).to(device)
-    with torch.no_grad():
-        lp = float(lpips_model(L1, L2))
-
-    # BERTScore
-    if cap1 and cap2:
-        _, _, F = score([cap1],[cap2], lang="en", verbose=False)
-        bert_f1 = float(F.mean())
-    else:
-        bert_f1 = 0.0
-
-    return clip_sim, lp, bert_f1
-
-def answer_vqa(question, image):
-    if not image or not question.strip():
-        return "Provide image + question."
-    try:
-        inputs_raw = vqa_processor(images=image, text=question, return_tensors="pt")
-        inputs = {k:v.to("cpu") for k,v in inputs_raw.items()}
-        with torch.no_grad():
-            out = vqa_model(**inputs)
-        ans_id = out.logits.argmax(-1)
-        return vqa_processor.decode(ans_id[0], skip_special_tokens=True)
-    except:
-        return "I could not determine the answer."
-
-# ==============================
-# Gradio UI
-# ==============================
-def build_ui():
-    with gr.Blocks(title="Multimodal AI Image Studio") as demo:
-        images_state = gr.State([None, None, None])
-        captions_state = gr.State(["", "", ""])
-
-        gr.Markdown("## Multimodal AI Image Studio (HF-ready)")
-
-        # --- Step 1: Upload Reference ---
-        upload_input = gr.Image(label="Upload Reference Image", type="pil")
-        upload_btn = gr.Button("Upload & Caption")
-        upload_preview = gr.Image(interactive=False)
-        caption_out = gr.Markdown()
-
-        def upload_and_caption(img, images_state, captions_state):
-            if img is None:
-                return None, "No image uploaded.", images_state, captions_state
-            caption = caption_for_image(img)
-            images_state[0] = img
-            captions_state[0] = caption
-            return img, caption, images_state, captions_state
-
-        upload_btn.click(upload_and_caption, inputs=[upload_input, images_state, captions_state],
-                         outputs=[upload_preview, caption_out, images_state, captions_state])
-
-        # --- Step 2: Generate SDXL & DreamShaper ---
-        sd_btn = gr.Button("Generate SD-Turbo")
-        ds_btn = gr.Button("Generate DreamShaper")
-        sd_preview = gr.Image(interactive=False)
-        ds_preview = gr.Image(interactive=False)
-
-        def gen_sd(caption, images_state, captions_state):
-            img = generate_image(gen_pipe, caption, enhancer="", negative="", seed=42, style="Photorealistic")
-            if img:
-                images_state[1] = img
-                captions_state[1] = caption_for_image(img)
-            return img, images_state, captions_state
-
-        def gen_ds(caption, images_state, captions_state):
-            img = generate_image(dreamshaper_pipe, caption, enhancer="", negative="", seed=123, style="Photorealistic")
-            if img:
-                images_state[2] = img
-                captions_state[2] = caption_for_image(img)
-            return img, images_state, captions_state
-
-        sd_btn.click(gen_sd, inputs=[caption_out, images_state, captions_state],
-                     outputs=[sd_preview, images_state, captions_state])
-        ds_btn.click(gen_ds, inputs=[caption_out, images_state, captions_state],
-                     outputs=[ds_preview, images_state, captions_state])
-
-        # --- Step 3: Metrics ---
-        metrics_btn = gr.Button("Compute Metrics")
-        metrics_out = gr.Markdown()
-
-        def metrics_ui(images_state, captions_state):
-            imgs = images_state or []
-            caps = captions_state or []
-            if None in imgs or "" in caps:
-                return "All three images and captions are required."
-            A = compute_metrics(imgs, caps, 0, 1)
-            B = compute_metrics(imgs, caps, 0, 2)
-            C = compute_metrics(imgs, caps, 1, 2)
-            return f"Reference ↔ SD-Turbo: {A}\nReference ↔ DreamShaper: {B}\nSD-Turbo ↔ DreamShaper: {C}"
-
-        metrics_btn.click(metrics_ui, inputs=[images_state, captions_state], outputs=[metrics_out])
-
-        # --- Step 4: NLP ---
-        nlp_btn = gr.Button("Analyze Captions")
-        nlp_out = gr.HTML()
-
-        def analyze_nlp(captions_state):
-            caps = captions_state or []
-            if "" in caps:
-                return "<b>All three captions are required.</b>"
-            labels = ["Reference", "SD-Turbo", "DreamShaper"]
-            html_blocks = []
-            for label, cap in zip(labels, caps):
-                # Sentiment
-                sentiment = "<br>".join([f"{s['label']}: {s['score']:.2f}" for s in sentiment_model(cap)])
-                # Entities
-                ents_list = ner_model(cap)
-                ents = "<br>".join([f"{e['entity_group']}: {e['word']}" for e in ents_list])
-                # Topics
-                topics_data = topic_model(cap, candidate_labels=['people','animals','objects','food','nature'])
-                topics = "<br>".join([f"{l}: {sc:.2f}" for l, sc in zip(topics_data['labels'], topics_data['scores'])])
-                html_blocks.append(f"<div style='padding:10px;'><h3>{label}</h3><b>Sentiment</b><br>{sentiment}<br><b>Entities</b><br>{ents}<br><b>Topics</b><br>{topics}</div>")
-            return "<div style='display:flex;gap:20px;'>" + "".join(html_blocks) + "</div>"
-
-        nlp_btn.click(analyze_nlp, inputs=[captions_state], outputs=[nlp_out])
-
-        # --- Step 5: VQA ---
-        vqa_input = gr.Textbox(label="Ask about reference image")
-        vqa_btn = gr.Button("Get Answer")
-        vqa_out = gr.Markdown()
-
-        def vqa_ui(question, img):
-            return answer_vqa(question, img)
-
-        vqa_btn.click(vqa_ui, inputs=[vqa_input, upload_preview], outputs=[vqa_out])
-
-    return demo
-
-# Launch
-demo = build_ui()
-demo.launch()
-
-
-####################################################################################
-# ==============================
-# SECTION 1
-# ==============================
-
-# Libraries
+#Dumped code
+# =========================
+# LIBRARIES & DEVICE SETUP
+# =========================
 import torch
 import gradio as gr
 from PIL import Image
@@ -640,9 +302,10 @@ def free_gpu_cache():
     if device == "cuda":
         torch.cuda.empty_cache()
 
-# ==============================
+# =========================
 # MODELS
-# ==============================
+# =========================
+# Image generation
 gen_pipe = DiffusionPipeline.from_pretrained(
     "stabilityai/sdxl-turbo",
     torch_dtype=torch.float16 if device=="cuda" else torch.float32
@@ -653,6 +316,7 @@ dreamshaper_pipe = DiffusionPipeline.from_pretrained(
     torch_dtype=torch.float16 if device=="cuda" else torch.float32
 ).to(device)
 
+# Captioning
 captioner = pipeline(
     "image-to-text",
     model="Salesforce/blip-image-captioning-large",
@@ -660,6 +324,7 @@ captioner = pipeline(
     generate_kwargs={"max_new_tokens":256, "num_beams":5, "temperature":0.7}
 )
 
+# NLP
 sentiment_model = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english",
                            device=0 if device=="cuda" else -1)
 ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english",
@@ -667,13 +332,16 @@ ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-engl
 topic_model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli",
                        device=0 if device=="cuda" else -1)
 
+# VQA
 vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
 vqa_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to("cpu")
 
+# Metrics
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 lpips_model = lpips.LPIPS(net='alex').to(device)
 lpips_transform = T.Compose([T.ToTensor(), T.Resize((256,256))])
 
+# Styles
 style_map = {
     "Photorealistic": "photorealistic, ultra-detailed, 8k, cinematic lighting",
     "Real Life": "natural lighting, true-to-life colors, DSLR",
@@ -686,25 +354,21 @@ style_map = {
     "Macro": "macro lens shallow DOF",
     "Cyberpunk": "neon cyberpunk futuristic",
 }
-# Section 2
-# ==============================
-# SECTION 2 — FUNCTIONS
-# ==============================
+
+# =========================
+# IMAGE GENERATION FUNCTIONS
+# =========================
 def generate_image_with_enhancer(base_caption, enhancer, negative, seed, style, images):
     images = images or []
     base_caption = base_caption or ""
     enhancer = enhancer or ""
-
     final_prompt = f"{base_caption}, {enhancer}".strip(", ")
     final_prompt = f"{final_prompt}, {style_map.get(style,'')}".strip(", ")
-
     try:
         seed = int(seed)
     except:
         seed = 42
-
     generator = torch.Generator(device="cpu").manual_seed(seed)
-
     try:
         with torch.no_grad():
             out = gen_pipe(prompt=final_prompt, negative_prompt=negative, generator=generator)
@@ -712,10 +376,8 @@ def generate_image_with_enhancer(base_caption, enhancer, negative, seed, style,
     except Exception as e:
         print("SD Turbo failed:", e)
         img = None
-
     if img:
         images.append(img)
-
     free_gpu_cache()
     return img, images
 
@@ -723,17 +385,13 @@ def generate_dreamshaper_with_enhancer(base_caption, enhancer, negative, seed, s
     images = images or []
     base_caption = base_caption or ""
     enhancer = enhancer or ""
-
     final_prompt = f"{base_caption}, {enhancer}".strip(", ")
     final_prompt = f"{final_prompt}, {style_map.get(style,'')}".strip(", ")
-
     try:
         seed = int(seed)
     except:
         seed = 42
-
     generator = torch.Generator(device="cpu").manual_seed(seed)
-
     try:
         with torch.no_grad():
             out = dreamshaper_pipe(prompt=final_prompt, negative_prompt=negative, generator=generator)
@@ -741,13 +399,14 @@ def generate_dreamshaper_with_enhancer(base_caption, enhancer, negative, seed, s
     except Exception as e:
         print("DreamShaper failed:", e)
         img = None
-
     if img:
         images.append(img)
-
     free_gpu_cache()
     return img, images
 
+# =========================
+# CAPTIONING
+# =========================
 def caption_for_image(img):
     try:
         out = captioner(img)
@@ -755,6 +414,9 @@ def caption_for_image(img):
     except:
         return "Caption failed."
 
+# =========================
+# VQA
+# =========================
 def answer_vqa(question, image):
     if not image or not question.strip():
         return "Provide image + question."
@@ -766,8 +428,11 @@ def answer_vqa(question, image):
         ans_id = out.logits.argmax(-1)
         return vqa_processor.decode(ans_id[0], skip_special_tokens=True)
     except:
-        return "VQA failed."
+        return "I could not determine the answer."
 
+# =========================
+# METRICS
+# =========================
 def compute_metrics(images, captions, i1, i2):
     img1 = images[i1]
     img2 = images[i2]
@@ -797,190 +462,142 @@ def compute_metrics(images, captions, i1, i2):
 
     return clip_sim, lp, bert_f1
 
-# ---------------- Build Gradio UI with Custom Look (Fully Robust) ----------------
-def build_ui_with_custom_ui():
+# =========================
+# GRADIO UI BUILD
+# =========================
+def build_full_ui():
     with gr.Blocks(title="Multimodal AI Image Studio") as demo:
-        # ---------------- CSS Styling ----------------
+        # --- CSS Styling ---
         gr.HTML(
         <style>
         .heading-orange h2, .heading-orange h3 { color: #ff5500 !important; }
         .orange-btn button { background-color: #ff5500 !important; color: white !important; border-radius: 6px !important; height: 36px !important; font-weight: bold; }
         .teal-btn button { background-color: #008080 !important; color: white !important; border-radius: 6px !important; height: 40px !important; font-weight: bold; }
-
-        /* Horizontal thin spinner */
-        .loading-line {
-            height: 4px;
-            background: linear-gradient(90deg, #008080 0%, #00cccc 50%, #008080 100%);
-            background-size: 200% 100%;
-            animation: loading 1s linear infinite;
-        }
-        @keyframes loading {
-            0% { background-position: 200% 0; }
-            100% { background-position: -200% 0; }
-        }
-
-        /* Match enhancer box to upload button */
-        .enhancer-box textarea {
-            width: 100% !important;
-            height: 36px !important;
-            box-sizing: border-box;
-            font-size: 14px;
-        }
-
-        /* Equal-height styling for Step-1 columns */
-        .equal-height-row {
-            display: flex;
-            align-items: stretch;
-        }
-        .equal-height-row > .gr-column {
-            display: flex;
-            flex-direction: column;
-        }
+        .loading-line { height:4px; background: linear-gradient(90deg,#008080 0%,#00cccc 50%,#008080 100%); background-size: 200% 100%; animation: loading 1s linear infinite; }
+        @keyframes loading { 0% { background-position:200% 0; } 100% { background-position:-200% 0; } }
+        .enhancer-box textarea { width:100% !important; height:36px !important; box-sizing:border-box; font-size:14px; }
+        .equal-height-row { display:flex; align-items:stretch; }
+        .equal-height-row > .gr-column { display:flex; flex-direction:column; }
         </style>
         )
 
-        # ---------------- Heading ----------------
-        gr.Markdown("## Multimodal AI Image Studio: An Integrated Comparative Perspective", elem_classes="heading-orange")
-
-        # ---------------- States ----------------
-        images_state = gr.State([])
-        captions_state = gr.State([])
-
-        # ---------------- Step 1: Upload Reference Image ----------------
-        gr.Markdown("### Upload Reference Image", elem_classes="heading-orange")
+        # --- States ---
+        images_state = gr.State([None, None, None])
+        captions_state = gr.State(["", "", ""])
 
+        # =========================
+        # Section 1: Upload Reference Image
+        # =========================
+        gr.Markdown("## 1️⃣ Upload Reference Image", elem_classes="heading-orange")
         with gr.Row(elem_classes="equal-height-row"):
             with gr.Column(scale=1):
                 upload_input = gr.Image(label="Drag & Drop Image", type="pil")
                 upload_btn = gr.Button("Upload Image & Generate Caption", elem_classes="orange-btn")
+                enhancer_box = gr.Textbox(label="Prompt Enhancer (Optional)", placeholder="Example: 'at night with neon lights'", elem_classes="enhancer-box")
             with gr.Column(scale=1):
                 upload_preview = gr.Image(label="Uploaded Image", interactive=False)
-                enhancer_box = gr.Textbox(
-                    label="Add Prompt Enhancer (Optional)",
-                    placeholder="Example: 'at night with neon lights', 'wearing a red jacket', etc.",
-                    elem_classes="enhancer-box"
-                )
                 caption_out = gr.Markdown(label="Generated Caption")
 
-        # Safe caption generation
-        def upload_and_generate_caption_ui(img, images_state, captions_state):
+        # Upload & caption function
+        def upload_and_caption(img, images_state, captions_state):
             if img is None:
-                return None, "No image uploaded.", images_state or [], captions_state or []
-
-            images = [img]
+                return None, "No image uploaded.", images_state, captions_state
+            images_state[0] = img
             try:
-                output = captioner(img)
-                caption = output[0].get("generated_text", "Caption failed.") if output else "Caption failed."
-            except Exception as e:
-                print("Captioning error:", e)
-                caption = "Caption failed."
-
-            captions = [caption]
-            return img, caption, images, captions
-
-        upload_btn.click(
-            upload_and_generate_caption_ui,
-            inputs=[upload_input, images_state, captions_state],
-            outputs=[upload_preview, caption_out, images_state, captions_state]
-        )
+                cap = caption_for_image(img)
+            except:
+                cap = "Caption failed."
+            captions_state[0] = cap
+            return img, cap, images_state, captions_state
 
-        # ---------------- Step 2: Generate SD-Turbo & DreamShaper ----------------
-        gr.Markdown("### Generate Images from Caption", elem_classes="heading-orange")
+        upload_btn.click(upload_and_caption, inputs=[upload_input, images_state, captions_state],
+                         outputs=[upload_preview, caption_out, images_state, captions_state])
+
+        # =========================
+        # Section 2: Generate SD-Turbo & DreamShaper
+        # =========================
+        gr.Markdown("## 2️⃣ Generate Images from Caption", elem_classes="heading-orange")
         with gr.Row():
-            with gr.Column(scale=1, min_width=300):
+            with gr.Column(scale=1):
                 sd_btn = gr.Button("Generate SD-Turbo Image", elem_classes="orange-btn")
                 sd_preview = gr.Image(label="SD-Turbo Image", interactive=False)
-            with gr.Column(scale=1, min_width=300):
+            with gr.Column(scale=1):
                 ds_btn = gr.Button("Generate DreamShaper Image", elem_classes="orange-btn")
                 ds_preview = gr.Image(label="DreamShaper Image", interactive=False)
 
-        def generate_sd_from_caption_ui(caption, enhancer, images_state, captions_state):
-            images_state = images_state or []
-            captions_state = captions_state or []
-            img, images = generate_image_with_enhancer(caption, enhancer="", negative="", seed=42, style="Photorealistic", images=images_state)
+        # Generate SD-Turbo
+        def generate_sd(caption, enhancer, images_state, captions_state):
+            img, images_state = generate_image_with_enhancer(caption, enhancer, negative="", seed=42, style="Photorealistic", images=images_state)
             if img:
-                try:
-                    generated_caption = captioner(img)[0].get("generated_text", "Caption failed.")
-                except:
-                    generated_caption = "Caption failed."
-                if len(captions_state) >= 2:
-                    captions_state[1] = generated_caption
-                else:
-                    captions_state.append(generated_caption)
-            return img, images, captions_state
-
-        def generate_ds_from_caption_ui(caption, enhancer, images_state, captions_state):
-            images_state = images_state or []
-            captions_state = captions_state or []
-            img, images = generate_dreamshaper_with_enhancer(caption, enhancer="", negative="", seed=123, style="Photorealistic", images=images_state)
+                captions_state[1] = caption_for_image(img)
+            return img, images_state, captions_state
+
+        # Generate DreamShaper
+        def generate_ds(caption, enhancer, images_state, captions_state):
+            img, images_state = generate_dreamshaper_with_enhancer(caption, enhancer, negative="", seed=123, style="Photorealistic", images=images_state)
             if img:
-                try:
-                    generated_caption = captioner(img)[0].get("generated_text", "Caption failed.")
-                except:
-                    generated_caption = "Caption failed."
-                if len(captions_state) >= 3:
-                    captions_state[2] = generated_caption
-                else:
-                    captions_state.append(generated_caption)
-            return img, images, captions_state
-
-        sd_btn.click(generate_sd_from_caption_ui, inputs=[caption_out, enhancer_box, images_state, captions_state],
+                captions_state[2] = caption_for_image(img)
+            return img, images_state, captions_state
+
+        sd_btn.click(generate_sd, inputs=[caption_out, enhancer_box, images_state, captions_state],
                      outputs=[sd_preview, images_state, captions_state])
-        ds_btn.click(generate_ds_from_caption_ui, inputs=[caption_out, enhancer_box, images_state, captions_state],
+        ds_btn.click(generate_ds, inputs=[caption_out, enhancer_box, images_state, captions_state],
                      outputs=[ds_preview, images_state, captions_state])
 
-        # ---------------- Step 3: Compute Pairwise Metrics ----------------
-        gr.Markdown("### Compute Pairwise Metrics", elem_classes="heading-orange")
+        # =========================
+        # Section 3: Compute Pairwise Metrics (Side-by-Side)
+        # =========================
+        gr.Markdown("## 3️⃣ Compute Pairwise Metrics", elem_classes="heading-orange")
         metrics_btn = gr.Button("Compute Metrics for All Pairs", elem_classes="teal-btn")
-        with gr.Row():
-            metrics_spinner = gr.HTML("<div style='height:4px;'></div>")  # single spinner
-        with gr.Row():
-            metrics_A = gr.Markdown()
-            metrics_B = gr.Markdown()
-            metrics_C = gr.Markdown()
-
-        def compute_metrics_all_pairs_ui(images, captions):
-            images = images or []
-            captions = captions or []
-            # show spinner
-            yield "<div class='loading-line'></div>", "", "", ""
-            if len(images) < 3 or len(captions) < 3:
-                msg = "All three images and captions are required to compute metrics."
-                yield "", msg, msg, msg
+        metrics_spinner = gr.HTML("<div style='height:4px;'></div>")
+        metrics_out = gr.HTML()
+
+        def compute_metrics_ui(images, captions):
+            yield "<div class='loading-line'></div>", ""
+            if any(i is None for i in images):
+                yield "All three images and captions are required."
             else:
                 try:
                     A = compute_metrics(images, captions, 0, 1)
                     B = compute_metrics(images, captions, 0, 2)
                     C = compute_metrics(images, captions, 1, 2)
-                    # remove spinner, show results
-                    yield "", f"**Reference ↔ SD-Turbo**\n{A}", f"**Reference ↔ DreamShaper**\n{B}", f"**SD-Turbo ↔ DreamShaper**\n{C}"
+                    def fmt(m):
+                        return f"CLIP: {m[0]:.3f}<br>LPIPS: {m[1]:.3f}<br>BERTScore F1: {m[2]:.3f}"
+                    html = f"""
+                    #<div style='display:flex; gap:40px; justify-content:space-around;'>
+                    #    <div style='text-align:center;'><b>Metrics A</b><br>{fmt(A)}</div>
+                    #    <div style='text-align:center;'><b>Metrics B</b><br>{fmt(B)}</div>
+                    #    <div style='text-align:center;'><b>Metrics C</b><br>{fmt(C)}</div>
+                    #</div>
+                    """
+                    yield html
                 except Exception as e:
-                    print("Metrics computation error:", e)
-                    msg = "Failed to compute metrics."
-                    yield "", msg, msg, msg
+                    print("Metrics error:", e)
+                    yield "Failed to compute metrics."
 
-        metrics_btn.click(compute_metrics_all_pairs_ui, inputs=[images_state, captions_state],
-                          outputs=[metrics_spinner, metrics_A, metrics_B, metrics_C])
+        metrics_btn.click(compute_metrics_ui, inputs=[images_state, captions_state],
+                          outputs=[metrics_out])
 
-        # ---------------- Step 4: NLP Analysis ----------------
-        gr.Markdown("### NLP Analysis of Captions", elem_classes="heading-orange")
+        # =========================
+        # Section 4: NLP Analysis
+        # =========================
+        gr.Markdown("## 4️⃣ NLP Analysis of Captions", elem_classes="heading-orange")
         nlp_btn = gr.Button("Analyze Captions", elem_classes="teal-btn")
-        nlp_spinner = gr.HTML("<div style='height:4px;'></div>")  # single spinner
+        nlp_spinner = gr.HTML("<div style='height:4px;'></div>")
         nlp_out = gr.HTML()
 
-        def analyze_caption_pipeline_ui(captions):
-            captions = captions or []
+        def analyze_captions_ui(captions):
             yield "<div class='loading-line'></div>", ""
-            if len(captions) < 3:
-                yield "", "<b>All three captions are required for NLP analysis.</b>"
+            if any(c=="" for c in captions):
+                yield "<b>All three captions are required for NLP analysis.</b>"
             else:
-                labels = ["Reference Image", "SD-Turbo", "DreamShaper"]
+                labels = ["Reference", "SD-Turbo", "DreamShaper"]
                 blocks = []
                 for label, caption in zip(labels, captions):
                     try:
                         sentiment = "<br>".join([f"{s['label']}: {s['score']:.2f}" for s in sentiment_model(caption)])
                     except:
-                        sentiment = "Sentiment analysis failed."
+                        sentiment = "Sentiment failed."
                     try:
                         ents_list = ner_model(caption)
                         ents = "<br>".join([f"{e.get('entity_group','')}: {e.get('word','')}" for e in ents_list]) or "None"
@@ -990,37 +607,44 @@ def build_ui_with_custom_ui():
                         topics_data = topic_model(caption, candidate_labels=['people','animals','objects','food','nature'])
                         topics = "<br>".join([f"{l}: {sc:.2f}" for l, sc in zip(topics_data.get('labels',[]), topics_data.get('scores',[]))])
                     except:
-                        topics = "Topic modeling failed."
+                        topics = "Topics failed."
                     block = f"<div style='flex:1;padding:10px;min-width:250px;'><h3><u>{label}</u></h3><b>Sentiment</b><br>{sentiment}<br><br><b>Entities</b><br>{ents}<br><br><b>Topics</b><br>{topics}</div>"
                     blocks.append(block)
-                yield "", f"<div style='display:flex; gap:20px; justify-content:space-between;'>{''.join(blocks)}</div>"
+                yield f"<div style='display:flex; gap:20px; justify-content:space-between;'>{''.join(blocks)}</div>"
 
-        nlp_btn.click(analyze_caption_pipeline_ui, inputs=[captions_state], outputs=[nlp_spinner, nlp_out])
+        nlp_btn.click(analyze_captions_ui, inputs=[captions_state], outputs=[nlp_out])
 
-        # ---------------- Step 5: Visual Question Answering (VQA) ----------------
-        gr.Markdown("### Visual Question Answering (VQA)", elem_classes="heading-orange")
+        # =========================
+        # Section 5: Visual Question Answering
+        # =========================
+        gr.Markdown("## 5️⃣ Visual Question Answering (VQA)", elem_classes="heading-orange")
         with gr.Row():
             with gr.Column(scale=1):
                 vqa_input = gr.Textbox(label="Enter a question about the reference image")
                 vqa_btn = gr.Button("Get Answer", elem_classes="teal-btn")
             with gr.Column(scale=1):
-                vqa_spinner = gr.HTML("<div style='height:4px;'></div>")  # single spinner
+                vqa_spinner = gr.HTML("<div style='height:4px;'></div>")
                 vqa_out = gr.Markdown(label="VQA Output")
 
-        def answer_vqa_ui(question, image):
+        def vqa_ui(question, image):
             yield "<div class='loading-line'></div>", ""
-            try:
-                ans = answer_vqa(question, image)
-            except Exception as e:
-                print("VQA error:", e)
-                ans = "I could not determine the answer."
-            yield "", ans
+            if not question.strip() or image is None:
+                yield "Provide image + question."
+            else:
+                try:
+                    ans = answer_vqa(question, image)
+                    yield f"<b>Answer:</b> {ans}"
+                except Exception as e:
+                    print("VQA error:", e)
+                    yield "Could not determine the answer."
 
-        vqa_btn.click(answer_vqa_ui, inputs=[vqa_input, upload_preview], outputs=[vqa_spinner, vqa_out])
+        vqa_btn.click(vqa_ui, inputs=[vqa_input, upload_preview], outputs=[vqa_out])
 
     return demo
 
-# Launch the interface
-demo = build_ui_with_custom_ui()
+# Launch
+demo = build_full_ui()
 demo.launch()
+
+
 """