Update app.py

app.py

@@ -18,9 +18,6 @@ from transformers import (
     T5Tokenizer,
 )
 
-# -------------------------------------------------
-# Device & models
-# -------------------------------------------------
 device = torch.device("cpu")
 
 IMG_MODEL = "nlpconnect/vit-gpt2-image-captioning"
@@ -34,13 +31,12 @@ rewriter_tok = T5Tokenizer.from_pretrained(TXT_MODEL)
 rewriter = T5ForConditionalGeneration.from_pretrained(TXT_MODEL).to(device).eval()
 
 
-# -------------------------------------------------
-# Helpers
-# -------------------------------------------------
 def load_image(url: str):
     """Return (PIL.Image, None) or (None, error). Handles http/https and data‑URL."""
     try:
-        url = url.strip()
+        url = (url or "").strip()
+        if not url:
+            return None, "No URL provided."
         if url.startswith("data:"):
             _, data = url.split(",", 1)
             img = Image.open(BytesIO(base64.b64decode(data))).convert("RGB")
@@ -55,10 +51,8 @@ def load_image(url: str):
 
 
 def generate_base(img: Image.Image, max_len=40, beams=2, sample=False):
-    """Return the longest caption (most detailed) from the vision model."""
     inputs = processor(images=img, return_tensors="pt")
     pix = inputs.pixel_values.to(device)
-
     if sample:
         out = vision.generate(
             pix,
@@ -83,12 +77,10 @@ def generate_base(img: Image.Image, max_len=40, beams=2, sample=False):
 
 
 def expand_caption(base: str, prompt: str = None, max_len=160):
-    """Use T5 to expand the base caption."""
     if prompt and prompt.strip():
         instr = f"Expand using: '{prompt}'. Caption: \"{base}\""
     else:
         instr = f"Expand with rich visual detail. Caption: \"{base}\""
-
     toks = rewriter_tok(
         instr,
         return_tensors="pt",
@@ -96,7 +88,6 @@ def expand_caption(base: str, prompt: str = None, max_len=160):
         padding="max_length",
         max_length=256,
     ).to(device)
-
     out = rewriter.generate(
         **toks,
         max_length=max_len,
@@ -108,36 +99,26 @@ def expand_caption(base: str, prompt: str = None, max_len=160):
 
 
 def async_expand(base, prompt, max_len, status):
-    """Background expansion; updates status dict."""
     try:
         status["text"] = "Expanding…"
-        time.sleep(0.1)  # tiny yield for UI responsiveness
+        time.sleep(0.1)
         result = expand_caption(base, prompt, max_len)
         status["text"] = "Done"
-        return result
+        status["final"] = result
     except Exception as e:
         status["text"] = f"Error: {e}"
-        return base
+        status["final"] = base
 
 
-# -------------------------------------------------
-# Gradio callbacks
-# -------------------------------------------------
 def fast_describe(url, prompt, detail, beams, sample):
     img, err = load_image(url)
     if err:
         return None, "", err
-
     detail_map = {"Low": 80, "Medium": 140, "High": 220}
     max_expand = detail_map.get(detail, 140)
-
     base = generate_base(img, beams=beams, sample=sample)
-    status = {"text": "Queued…"}
-
-    def worker():
-        status["final"] = async_expand(base, prompt, max_expand, status)
-
-    threading.Thread(target=worker, daemon=True).start()
+    status = {"text": "Queued…", "final": ""}
+    threading.Thread(target=async_expand, args=(base, prompt, max_expand, status), daemon=True).start()
     return img, base, status["text"]
 
 
@@ -145,10 +126,8 @@ def final_caption(url, prompt, detail, beams, sample):
     img, err = load_image(url)
     if err:
         return "", err
-
     detail_map = {"Low": 80, "Medium": 140, "High": 220}
     max_expand = detail_map.get(detail, 140)
-
     base = generate_base(img, beams=beams, sample=sample)
     try:
         final = expand_caption(base, prompt, max_expand)
@@ -157,9 +136,6 @@ def final_caption(url, prompt, detail, beams, sample):
         return base, f"Expand error: {e}"
 
 
-# -------------------------------------------------
-# UI
-# -------------------------------------------------
 css = "footer {display:none !important;}"
 with gr.Blocks() as demo:
     gr.Markdown("## Image Describer")
@@ -189,15 +165,6 @@ with gr.Blocks() as demo:
         outputs=[caption_out, status_out],
     )
 
-# -------------------------------------------------
-# Launch
-# -------------------------------------------------
 if __name__ == "__main__":
     demo.queue()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        css=css,
-        title="Image Describer (CPU)",
-        prevent_thread_lock=True,
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, css=css, prevent_thread_lock=True)