Update app.py

app.py

@@ -1,7 +1,7 @@
+# app.py – minimal, CPU‑only, high‑quality captions
 import gradio as gr
 from PIL import Image
-import requests
-from io import BytesIO
+import requests, urllib.parse, threading, time
 import torch
 from transformers import (
     VisionEncoderDecoderModel,
@@ -10,195 +10,174 @@ from transformers import (
     T5ForConditionalGeneration,
     T5Tokenizer,
 )
-import urllib.parse
-import threading
-import time
 
+# -------------------------------------------------
+# Device & models (CPU)
+# -------------------------------------------------
 device = torch.device("cpu")
 
-# Model names
-PROCESSOR_NAME = "nlpconnect/vit-gpt2-image-captioning"
-REWRITER_NAME = "t5-small"
+IMG_MODEL = "nlpconnect/vit-gpt2-image-captioning"
+TXT_MODEL = "t5-small"
 
-# Load models (CPU)
-processor = ViTImageProcessor.from_pretrained(PROCESSOR_NAME)
-tokenizer = AutoTokenizer.from_pretrained(PROCESSOR_NAME)
-model = VisionEncoderDecoderModel.from_pretrained(PROCESSOR_NAME).to(device)
-model.eval()
+processor = ViTImageProcessor.from_pretrained(IMG_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(IMG_MODEL)
+vision = VisionEncoderDecoderModel.from_pretrained(IMG_MODEL).to(device).eval()
 
-rewriter_tokenizer = T5Tokenizer.from_pretrained(REWRITER_NAME)
-rewriter = T5ForConditionalGeneration.from_pretrained(REWRITER_NAME).to(device)
-rewriter.eval()
+rewriter_tok = T5Tokenizer.from_pretrained(TXT_MODEL)
+rewriter = T5ForConditionalGeneration.from_pretrained(TXT_MODEL).to(device).eval()
 
-def load_image_from_url(url: str, timeout=10):
+# -------------------------------------------------
+# Helpers
+# -------------------------------------------------
+def load_image(url: str):
+    """Return PIL image or (None, error). Handles http/https and data‑URL."""
     try:
         url = url.strip()
         if url.startswith("data:"):
-            header, encoded = url.split(",", 1)
             import base64
-            data = base64.b64decode(encoded)
-            img = Image.open(BytesIO(data)).convert("RGB")
+            _, data = url.split(",", 1)
+            img = Image.open(BytesIO(base64.b64decode(data))).convert("RGB")
             return img, None
-        parsed = urllib.parse.urlsplit(url)
-        if parsed.scheme == "":
-            return None, "Invalid URL (missing scheme: http/https)."
-        resp = requests.get(url, timeout=timeout, headers={"User-Agent": "huggingface-space/1.0"})
-        resp.raise_for_status()
-        img = Image.open(BytesIO(resp.content)).convert("RGB")
-        return img, None
+        if not urllib.parse.urlsplit(url).scheme:
+            return None, "Missing http/https scheme."
+        r = requests.get(url, timeout=10, headers={"User-Agent": "duck.ai"})
+        r.raise_for_status()
+        return Image.open(BytesIO(r.content)).convert("RGB"), None
     except Exception as e:
-        return None, f"Error loading image: {e}"
+        return None, f"Load error: {e}"
 
-# --- Generation & rewriting helpers ---
-def generate_caption_candidates(img: Image.Image, max_len: int = 40, num_beams: int = 2, num_return_sequences: int = 3, do_sample: bool = False):
+def generate_base(img: Image.Image, max_len=40, beams=2, sample=False):
+    """Return a single “most detailed” base caption."""
     inputs = processor(images=img, return_tensors="pt")
-    pixel_values = inputs.pixel_values.to(device)
-
-    gen_kwargs = {
-        "max_length": max_len,
-        "num_beams": num_beams,
-        "early_stopping": True,
-        "do_sample": do_sample,
-        "num_return_sequences": num_return_sequences,
-    }
-    # model.generate returns tensor of shape (num_return_sequences, seq_len) when requested
-    outputs = model.generate(pixel_values, **gen_kwargs)
-    captions = [tokenizer.decode(o, skip_special_tokens=True).strip() for o in outputs]
-    # Deduplicate preserving order
-    seen = set()
-    unique = []
-    for c in captions:
-        if c not in seen:
-            seen.add(c)
-            unique.append(c)
-    return unique
-
-def pick_most_detailed(candidates):
-    # heuristic: prefer longer by word count, then more unique words
-    best = max(candidates, key=lambda s: (len(s.split()), len(set(s.split()))))
-    return best
-
-def expand_with_t5(caption: str, prompt: str = None, max_len: int = 160):
-    # Instruction to expand and add rich visual detail
+    pix = inputs.pixel_values.to(device)
+
+    if sample:
+        out = vision.generate(
+            pix,
+            max_length=max_len,
+            do_sample=True,
+            temperature=0.8,
+            top_k=50,
+            top_p=0.9,
+            num_return_sequences=3,
+            early_stopping=True,
+        )
+    else:
+        # ensure num_return ≤ beams
+        out = vision.generate(
+            pix,
+            max_length=max_len,
+            num_beams=beams,
+            num_return_sequences=min(3, beams),
+            early_stopping=True,
+        )
+    caps = [tokenizer.decode(o, skip_special_tokens=True).strip() for o in out]
+    # pick longest (most detailed)
+    return max(caps, key=lambda s: len(s.split()))
+
+def expand_caption(base: str, prompt: str = None, max_len=160):
+    """Rich T5 expansion."""
     if prompt and prompt.strip():
-        instr = f"Expand and elaborate the caption using this instruction: '{prompt}'. Caption: \"{caption}\""
+        instr = f"Expand using: '{prompt}'. Caption: \"{base}\""
     else:
-        instr = f"Expand and elaborate the caption with rich visual detail (objects, colors, textures, scene, actions). Caption: \"{caption}\""
-    tok = rewriter_tokenizer(instr, return_tensors="pt", truncation=True, padding=True).to(device)
-    out = rewriter.generate(**tok, max_length=max_len, num_beams=4, early_stopping=True, no_repeat_ngram_size=3)
-    expanded = rewriter_tokenizer.decode(out[0], skip_special_tokens=True).strip()
-    return expanded
-
-# Background worker pattern to run expansion and report progress
-def _background_expand_and_return(caption, prompt, max_expand_len, status_callback):
+        instr = f"Expand with rich visual detail. Caption: \"{base}\""
+
+    toks = rewriter_tok(
+        instr,
+        return_tensors="pt",
+        truncation=True,
+        padding="max_length",
+        max_length=256,
+    ).to(device)
+
+    out = rewriter.generate(
+        **toks,
+        max_length=max_len,
+        num_beams=4,
+        early_stopping=True,
+        no_repeat_ngram_size=3,
+    )
+    return rewriter_tok.decode(out[0], skip_special_tokens=True).strip()
+
+# -------------------------------------------------
+# Async expansion (background thread)
+# -------------------------------------------------
+def async_expand(base, prompt, max_len, status):
     try:
-        # Inform start
-        status_callback("Expanding caption (step 1/2)...")
-        # Small sleep allows UI update
-        time.sleep(0.1)
-        expanded = expand_with_t5(caption, prompt=prompt, max_len=max_expand_len)
-        status_callback("Finalizing (step 2/2)...")
+        status["text"] = "Expanding…"
         time.sleep(0.1)
-        status_callback("Done")
-        return expanded
+        result = expand_caption(base, prompt, max_len)
+        status["text"] = "Done"
+        return result
     except Exception as e:
-        status_callback(f"Error during expand: {e}")
-        return caption
-
-# Main describe function used by Gradio; it triggers generation and then expansion in background
-def describe_image_controller(url: str, prompt: str, detail_level: str, max_caption_len: int = 40, beams: int = 2, do_sample: bool = True):
-    """
-    Returns: (img or None, caption_text, status_text)
-    The UI will start background expansion and update status via a small helper.
-    """
-    img, err = load_image_from_url(url)
+        status["text"] = f"Error: {e}"
+        return base
+
+# -------------------------------------------------
+# Gradio callbacks
+# -------------------------------------------------
+def fast_describe(url, prompt, detail, beams, sample):
+    img, err = load_image(url)
     if err:
-        return None, "", f"Error: {err}"
+        return None, "", err
 
-    # Map detail_level to rewriter max_len
     detail_map = {"Low": 80, "Medium": 140, "High": 220}
-    max_expand_len = detail_map.get(detail_level, 140)
-
-    # Generate candidates
-    candidates = generate_caption_candidates(img, max_len=max_caption_len, num_beams=beams, num_return_sequences=3, do_sample=do_sample)
-    base = pick_most_detailed(candidates)
+    max_expand = detail_map.get(detail, 140)
 
-    # Start background thread to expand (T5) and update status via a Gradio status element (we'll use a simple polling text)
-    # We'll use a small mutable container to send status updates via closure
-    status = {"text": "Queued for expansion..."}
-    def status_callback(s):
-        status["text"] = s
-
-    result_container = {"final": base}
+    base = generate_base(img, beams=beams, sample=sample)
+    status = {"text": "Queued…"}
 
     def worker():
-        expanded = _background_expand_and_return(base, prompt, max_expand_len, status_callback)
-        result_container["final"] = expanded
-
-    thread = threading.Thread(target=worker, daemon=True)
-    thread.start()
+        status["final"] = async_expand(base, prompt, max_expand, status)
 
-    # Return image, initial base caption, and initial status. The frontend will poll for status/final via separate endpoints
+    threading.Thread(target=worker, daemon=True).start()
     return img, base, status["text"]
 
-# Polling endpoints to retrieve status and final caption
-def poll_status_and_caption(url: str, prompt: str, _placeholder):
-    # In this simple pattern we re-run a lightweight check by storing results in a global map keyed by URL+prompt
-    # For simplicity in this Space we will re-run expansion synchronously here if needed.
-    # But to avoid redoing heavy work, you can implement a shared cache (omitted for brevity).
-    return "If expansion still running, refresh in a few seconds. Final caption will replace base when ready."
-
-# Simple endpoint to get final expanded caption synchronously (used when user hits 'Get final caption')
-def get_final_caption(url: str, prompt: str, detail_level: str, max_caption_len: int = 40, beams: int = 2, do_sample: bool = True):
-    img, err = load_image_from_url(url)
+def final_caption(url, prompt, detail, beams, sample):
+    img, err = load_image(url)
     if err:
-        return "", f"Error: {err}"
-    candidates = generate_caption_candidates(img, max_len=max_caption_len, num_beams=beams, num_return_sequences=3, do_sample=do_sample)
-    base = pick_most_detailed(candidates)
+        return "", err
     detail_map = {"Low": 80, "Medium": 140, "High": 220}
-    max_expand_len = detail_map.get(detail_level, 140)
+    max_expand = detail_map.get(detail, 140)
+
+    base = generate_base(img, beams=beams, sample=sample)
     try:
-        expanded = expand_with_t5(base, prompt=prompt, max_len=max_expand_len)
-        return expanded, "Done"
+        final = expand_caption(base, prompt, max_expand)
+        return final, "Done"
     except Exception as e:
         return base, f"Expand error: {e}"
 
-# Gradio UI
-css = """
-footer {display: none !important;}
-"""
-
-with gr.Blocks(css=css, title="Image Describer (vit-gpt2, promptable, detailed)") as demo:
-    gr.Markdown("## Image Describer — uncensored captions, optional prompt to bias description. Use 'Get final caption' for the detailed expanded output (may take longer).")
+# -------------------------------------------------
+# UI
+# -------------------------------------------------
+css = "footer {display:none !important;}"
+with gr.Blocks(css=css, title="Image Describer (CPU)") as demo:
+    gr.Markdown("## Image Describer – fast base caption + optional detailed rewrite")
     with gr.Row():
-        with gr.Column(scale=1):
-            url_in = gr.Textbox(label="Image URL or data URL", placeholder="https://example.com/photo.jpg")
-            prompt_in = gr.Textbox(label="Optional prompt (e.g. 'Focus on people and actions')", placeholder="Focus on people, actions, or colors.")
-            detail_level = gr.Radio(choices=["Low", "Medium", "High"], value="Medium", label="Detail level (affects expansion length)")
-            max_len = gr.Slider(minimum=8, maximum=80, value=40, label="Base caption max length")
-            beams = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Num beams (higher = better quality, slower)")
-            do_sample_chk = gr.Checkbox(label="Enable sampling (more diverse)", value=True)
-            go = gr.Button("Load & Describe (fast)")
-            get_final = gr.Button("Get final caption (detailed, slower)")
-            status_txt = gr.Textbox(label="Status", value="Idle", interactive=False)
-        with gr.Column(scale=1):
+        with gr.Column():
+            url_in = gr.Textbox(label="Image URL / data‑URL")
+            prompt_in = gr.Textbox(label="Optional prompt")
+            detail_in = gr.Radio(["Low", "Medium", "High"], value="Medium", label="Detail level")
+            beams_in = gr.Slider(1, 4, step=1, value=2, label="Beams (higher = better, slower)")
+            sample_in = gr.Checkbox(label="Enable sampling (more diverse)", value=False)
+            go_btn = gr.Button("Load & Describe (fast)")
+            final_btn = gr.Button("Get final caption (detailed)")
+            status_out = gr.Textbox(label="Status", interactive=False)
+        with gr.Column():
             img_out = gr.Image(type="pil", label="Image")
-        with gr.Column(scale=1):
-            caption_out = gr.Textbox(label="Caption (base or final)", lines=8)
-
-    # Fast path: generate base caption and immediately start background expand (status will be approximate)
-    def on_go(url, prompt, detail_level, max_len, beams, do_sample):
-        img, base_caption, status = describe_image_controller(url, prompt, detail_level, max_caption_len=max_len, beams=beams, do_sample=do_sample)
-        return img, base_caption, status
-
-    go.click(fn=on_go, inputs=[url_in, prompt_in, detail_level, max_len, beams, do_sample_chk], outputs=[img_out, caption_out, status_txt])
-
-    # Synchronous, explicit final result (user clicks when they want the full expanded caption)
-    def on_get_final(url, prompt, detail_level, max_len, beams, do_sample):
-        final_caption, status = get_final_caption(url, prompt, detail_level, max_caption_len=max_len, beams=beams, do_sample=do_sample)
-        return final_caption, status
-
-    get_final.click(fn=on_get_final, inputs=[url_in, prompt_in, detail_level, max_len, beams, do_sample_chk], outputs=[caption_out, status_txt])
+        with gr.Column():
+            caption_out = gr.Textbox(label="Caption", lines=8)
+
+    go_btn.click(
+        fn=fast_describe,
+        inputs=[url_in, prompt_in, detail_in, beams_in, sample_in],
+        outputs=[img_out, caption_out, status_out],
+    )
+    final_btn.click(
+        fn=final_caption,
+        inputs=[url_in, prompt_in, detail_in, beams_in, sample_in],
+        outputs=[caption_out, status_out],
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)