Update app.py

app.py

@@ -11,33 +11,34 @@ from transformers import (
     T5Tokenizer,
 )
 import urllib.parse
+import threading
+import time
 
 device = torch.device("cpu")
 
-# Models
+# Model names
 PROCESSOR_NAME = "nlpconnect/vit-gpt2-image-captioning"
+REWRITER_NAME = "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()
 
-# Optional rewriter (T5-small) to make captions more natural / respond to prompt
-rewriter_tokenizer = T5Tokenizer.from_pretrained("t5-small")
-rewriter = T5ForConditionalGeneration.from_pretrained("t5-small").to(device)
+rewriter_tokenizer = T5Tokenizer.from_pretrained(REWRITER_NAME)
+rewriter = T5ForConditionalGeneration.from_pretrained(REWRITER_NAME).to(device)
 rewriter.eval()
 
 def load_image_from_url(url: str, timeout=10):
     try:
-        # allow string that is data URL or direct URL
         url = url.strip()
         if url.startswith("data:"):
-            # let PIL handle data URLs via BytesIO after splitting
             header, encoded = url.split(",", 1)
             import base64
             data = base64.b64decode(encoded)
             img = Image.open(BytesIO(data)).convert("RGB")
             return img, None
-        # ensure proper URL encoding
         parsed = urllib.parse.urlsplit(url)
         if parsed.scheme == "":
             return None, "Invalid URL (missing scheme: http/https)."
@@ -48,68 +49,156 @@ def load_image_from_url(url: str, timeout=10):
     except Exception as e:
         return None, f"Error loading image: {e}"
 
-def generate_caption(img: Image.Image, prompt: str = None, max_len: int = 30, num_beams: int = 2):
-    # Prepare encoder inputs
+# --- 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):
     inputs = processor(images=img, return_tensors="pt")
     pixel_values = inputs.pixel_values.to(device)
 
-    # If a prompt is provided, prepend it to the decoder start tokens via tokenizer (prefix)
-    # This is a lightweight way to bias output by using the tokenizer's bos/tokenizer decoding prefix.
-    gen_kwargs = {"max_length": max_len, "num_beams": num_beams, "early_stopping": True}
-    if prompt:
-        # For vit-gpt2 model, we can try to use forced_decoder_input_ids or prefix decoding
-        # Simpler approach: generate normally and then rely on rewriter to apply prompt.
-        pass
-
-    out = model.generate(pixel_values, **gen_kwargs)
-    caption = tokenizer.decode(out[0], skip_special_tokens=True).strip()
-    return caption
-
-def rewrite_caption_with_prompt(caption: str, prompt: str = None, max_len: int = 64):
-    # If prompt provided, use it to instruct T5; otherwise paraphrase
-    if prompt:
-        input_text = f"paraphrase: {caption} prompt: {prompt}"
+    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
+    if prompt and prompt.strip():
+        instr = f"Expand and elaborate the caption using this instruction: '{prompt}'. Caption: \"{caption}\""
     else:
-        input_text = "paraphrase: " + caption
-    tok = rewriter_tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
-    out = rewriter.generate(**tok, max_length=max_len, num_beams=2, early_stopping=True)
-    rewritten = rewriter_tokenizer.decode(out[0], skip_special_tokens=True).strip()
-    return rewritten
+        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):
+    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)...")
+        time.sleep(0.1)
+        status_callback("Done")
+        return expanded
+    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)
+    if err:
+        return None, "", f"Error: {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)
 
-def describe_image(url: str, prompt: str, max_caption_len: int = 30, expand: bool = True, beams: int = 2):
+    # 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}
+
+    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()
+
+    # Return image, initial base caption, and initial status. The frontend will poll for status/final via separate endpoints
+    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)
     if err:
-        return None, f"Error: {err}"
-    caption = generate_caption(img, prompt=prompt, max_len=max_caption_len, num_beams=beams)
-    if expand:
-        try:
-            caption = rewrite_caption_with_prompt(caption, prompt=prompt, max_len=64)
-        except Exception:
-            pass
-    if len(caption.split()) < 6:
-        caption = f"{caption}. The scene appears to contain: {caption.lower()}."
-    return img, caption
+        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)
+    detail_map = {"Low": 80, "Medium": 140, "High": 220}
+    max_expand_len = detail_map.get(detail_level, 140)
+    try:
+        expanded = expand_with_t5(base, prompt=prompt, max_len=max_expand_len)
+        return expanded, "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, uncensored, promptable)") as demo:
-    gr.Markdown("## Image Describer — uncensored captions, optional prompt to bias description")
+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).")
     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. 'Describe people and actions')", placeholder="Focus on people, actions, or colors.")
-            max_len = gr.Slider(minimum=8, maximum=60, value=30, label="Max caption length")
+            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)")
-            expand_chk = gr.Checkbox(label="Rewrite/Paraphrase with prompt (slower)", value=True)
-            go = gr.Button("Load & Describe")
+            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):
             img_out = gr.Image(type="pil", label="Image")
         with gr.Column(scale=1):
-            caption_out = gr.Textbox(label="Descriptive caption", lines=6)
+            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
 
-    go.click(fn=describe_image, inputs=[url_in, prompt_in, max_len, expand_chk, beams], outputs=[img_out, caption_out])
+    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])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)