Update app.py

app.py

@@ -1,95 +1,145 @@
 # app.py
-# VisPrompt Studio — Image/Text → Detailed Prompts (CPU)
-# Focus: ensure LONG, DETAILED scene descriptions (esp. people/adult photos)
-# Models: Qwen2-VL-2B-Instruct (default, richer captions) and BLIP base (faster)
-
-import os, random, math
+# VisPrompt Studio — Image/Text → Long, Detailed Prompts
+# - Runs on Hugging Face Spaces (free CPU)
+# - Stronger scene descriptions (esp. people/adult) using Qwen2-VL-2B-Instruct
+# - BLIP base as a fast fallback
+# - Model selector is non-typable (Radio) + server-side guard + automatic fallback
+# - Uses Textboxes (not Code) for copy-friendly outputs (Gradio 5.x safe)
+
+import os
+import math
+import random
 from typing import Optional
 
 import gradio as gr
 from PIL import Image, ExifTags
 import torch
 from transformers import (
-    AutoProcessor, AutoModelForVision2Seq,
-    BlipProcessor, BlipForConditionalGeneration,
+    AutoProcessor,
+    AutoModelForVision2Seq,
+    BlipProcessor,
+    BlipForConditionalGeneration,
 )
 
+# -------------------- Models --------------------
 MODEL_QWEN2_VL = "Qwen/Qwen2-VL-2B-Instruct"
-MODEL_BLIP     = "Salesforce/blip-image-captioning-base"
+MODEL_BLIP = "Salesforce/blip-image-captioning-base"
+VALID_MODELS = ("Qwen2-VL-2B-Instruct", "BLIP base")
 
+# Lazy globals
 _q_proc = _q_model = None
 _b_proc = _b_model = None
 
-def device():
-    return "cpu"  # (If you upgrade to ZeroGPU, switch to "cuda" when available)
 
-def safe_open_image(inp):
-    if inp is None: return None
-    if isinstance(inp, Image.Image): return inp.convert("RGB")
+def device() -> str:
+    """
+    Return runtime device. Keep 'cpu' for free Spaces.
+    If you later use ZeroGPU, you can switch to:
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    """
+    return "cpu"
+
+
+# -------------------- Helpers --------------------
+def safe_open_image(inp) -> Optional[Image.Image]:
+    if inp is None:
+        return None
+    if isinstance(inp, Image.Image):
+        return inp.convert("RGB")
     return Image.open(inp).convert("RGB")
 
+
 def exif_info(img: Image.Image) -> dict:
     out = {}
     try:
         raw = img.getexif()
-        if raw:
-            tag_map = {ExifTags.TAGS.get(k, k): v for k, v in raw.items()}
-            fl = tag_map.get("FocalLength")
-            if isinstance(fl, tuple) and len(fl)==2 and fl[1]!=0:
-                out["focal_length_mm"] = round(fl[0]/fl[1], 1)
-            elif isinstance(fl, (int,float)):
-                out["focal_length_mm"] = float(fl)
-            if "Model" in tag_map: out["camera_model"] = str(tag_map["Model"]).strip()
-            if "Make"  in tag_map: out["camera_make"]  = str(tag_map["Make"]).strip()
+        if not raw:
+            return out
+        tag_map = {ExifTags.TAGS.get(k, k): v for k, v in raw.items()}
+
+        fl = tag_map.get("FocalLength")
+        if isinstance(fl, tuple) and len(fl) == 2 and fl[1] != 0:
+            out["focal_length_mm"] = round(fl[0] / fl[1], 1)
+        elif isinstance(fl, (int, float)):
+            out["focal_length_mm"] = float(fl)
+
+        if "Model" in tag_map:
+            out["camera_model"] = str(tag_map["Model"]).strip()
+        if "Make" in tag_map:
+            out["camera_make"] = str(tag_map["Make"]).strip()
     except Exception:
         pass
     return out
 
-CAMERA_ANGLES = ["eye-level","low-angle","high-angle","aerial/top-down","worm's-eye","Dutch tilt"]
-SHOT_SIZES    = ["extreme close-up","close-up","medium shot","wide shot","extreme wide shot"]
-LENSES        = ["14mm ultra-wide","24mm wide-angle","35mm street","50mm standard","85mm portrait","135mm telephoto","200mm telephoto"]
-LIGHTING      = ["soft diffused light","golden hour rim light","overcast softbox feel","hard noon sun with crisp shadows","blue hour ambient","tungsten practicals","neon key with colored rim","studio softbox left + kicker right","volumetric god rays","moody chiaroscuro"]
-COMPOSITION   = ["rule of thirds","leading lines","central composition","frame within a frame","negative space","foreground bokeh","shallow depth of field","symmetry","diagonal dynamics","layered depth"]
-POST          = ["high micro-contrast","subtle film grain","cinematic color grade","Kodak Portra palette","teal–orange split tone","matte shadows","HDR tonality restrained","8k detail","photorealistic finish"]
-STYLES        = ["documentary realism","editorial fashion","street photography","nature landscape","architectural minimalism","cyberpunk neon","sci-fi cinematic","fantasy epic","moody noir","retro 35mm film"]
-VIDEO_MOTION  = ["slow dolly-in","gentle push-in","handheld micro-jitters","locked-off tripod","gimbal glide","orbit right-to-left","crane rise","parallax truck shot","rack focus mid-shot"]
-VIDEO_META    = ["24 fps","30 fps","48 fps","cinematic motion blur","shutter 1/48","global illumination feel"]
-VIDEO_TRANS   = ["straight cut","cinematic fade-in/out","cross dissolve","match cut"]
-NEGATIVE      = ["no distortion","no chromatic aberration","no blown highlights","no excessive sharpening","no artifacts","no watermark","no text overlay"]
+
+# Lexicons for enrichment
+CAMERA_ANGLES = ["eye-level", "low-angle", "high-angle", "aerial/top-down", "worm's-eye", "Dutch tilt"]
+SHOT_SIZES = ["extreme close-up", "close-up", "medium shot", "wide shot", "extreme wide shot"]
+LENSES = ["14mm ultra-wide", "24mm wide-angle", "35mm street", "50mm standard", "85mm portrait", "135mm telephoto", "200mm telephoto"]
+LIGHTING = [
+    "soft diffused light", "golden hour rim light", "overcast softbox feel", "hard noon sun with crisp shadows",
+    "blue hour ambient", "tungsten practicals", "neon key with colored rim", "studio softbox left + kicker right",
+    "volumetric god rays", "moody chiaroscuro"
+]
+COMPOSITION = [
+    "rule of thirds", "leading lines", "central composition", "frame within a frame", "negative space",
+    "foreground bokeh", "shallow depth of field", "symmetry", "diagonal dynamics", "layered depth"
+]
+POST = [
+    "high micro-contrast", "subtle film grain", "cinematic color grade", "Kodak Portra palette",
+    "teal–orange split tone", "matte shadows", "HDR tonality restrained", "8k detail", "photorealistic finish"
+]
+STYLES = [
+    "documentary realism", "editorial fashion", "street photography", "nature landscape", "architectural minimalism",
+    "cyberpunk neon", "sci-fi cinematic", "fantasy epic", "moody noir", "retro 35mm film"
+]
+VIDEO_MOTION = [
+    "slow dolly-in", "gentle push-in", "handheld micro-jitters", "locked-off tripod", "gimbal glide",
+    "orbit right-to-left", "crane rise", "parallax truck shot", "rack focus mid-shot"
+]
+VIDEO_META = ["24 fps", "30 fps", "48 fps", "cinematic motion blur", "shutter 1/48", "global illumination feel"]
+VIDEO_TRANSITIONS = ["straight cut", "cinematic fade-in/out", "cross dissolve", "match cut"]
+NEGATIVE = ["no distortion", "no chromatic aberration", "no blown highlights", "no excessive sharpening", "no artifacts", "no watermark", "no text overlay"]
+
 
 def smart_choices(caption: str, w: int, h: int, exif: dict, seed: int):
-    rnd    = random.Random(seed)
+    rnd = random.Random(seed)
     aspect = (w / h) if h > 0 else 1.0
-    shot   = "wide shot" if aspect >= 1.5 else "medium shot"
-    text   = caption.lower()
-    if any(k in text for k in ["portrait","selfie","face","headshot"]): shot = "close-up"
-    if any(k in text for k in ["aerial","drone","top-down","overhead","bird's-eye","bird’s-eye"]):
+
+    shot = "wide shot" if aspect >= 1.5 else "medium shot"
+    t = caption.lower()
+    if any(k in t for k in ["portrait", "selfie", "face", "headshot"]):
+        shot = "close-up"
+
+    if any(k in t for k in ["aerial", "drone", "top-down", "overhead", "bird's-eye", "bird’s-eye"]):
         angle = "aerial/top-down"
-    elif any(k in text for k in ["towering","imposing","hero","statue"]):
+    elif any(k in t for k in ["towering", "imposing", "hero", "statue"]):
         angle = "low-angle"
     else:
-        angle = rnd.choice(CAMERA_ANGLES[:3])
+        angle = rnd.choice(CAMERA_ANGLES[:3])  # eye/high/low
+
     lens = f'{exif["focal_length_mm"]}mm' if "focal_length_mm" in exif else rnd.choice(LENSES)
-    light= rnd.choice(LIGHTING)
+    light = ", ".join(rnd.sample(LIGHTING, k=2))
     comp = ", ".join(rnd.sample(COMPOSITION, k=3))
-    style= rnd.choice(STYLES)
+    style = rnd.choice(STYLES)
     post = ", ".join(rnd.sample(POST, k=3))
-    neg  = ", ".join(NEGATIVE[:3] + rnd.sample(NEGATIVE[3:], k=2))
+    neg = ", ".join(NEGATIVE[:3] + rnd.sample(NEGATIVE[3:], k=2))
     return angle, shot, lens, light, comp, style, post, neg, aspect
 
+
 def build_image_prompt(desc, angle, shot, lens, light, comp, style, post, aspect):
     ar = "16:9" if aspect >= 1.6 else ("4:3" if aspect >= 1.3 else "1:1")
-    # Expanded, sectioned prompt for richer output
     return (
         f"{desc}\n"
         f"Camera: {angle}; {shot}; lens {lens}.\n"
-        f"Lighting: {light}; natural falloff; realistic reflections and shadows.\n"
-        f"Composition: {comp}; controlled perspective; clean subject separation; coherent scale.\n"
+        f"Lighting: {light}; realistic reflections and shadow behavior.\n"
+        f"Composition: {comp}; clean subject separation; coherent scale and perspective.\n"
         f"Style: {style}; grounded materials and textures.\n"
         f"Post: {post}; minimal artifacts; print-ready clarity.\n"
         f"Aspect ratio hint: {ar}."
     )
 
+
 def build_video_prompt(img_prompt: str, motion: str, meta: list, transition: str, seconds: int):
     meta_str = ", ".join(meta)
     return (
@@ -98,95 +148,113 @@ def build_video_prompt(img_prompt: str, motion: str, meta: list, transition: str
         f"Maintain temporal consistency, realistic parallax, and stable lighting throughout."
     )
 
-# ---------------- Caption backends ----------------
+
+# -------------------- Caption backends --------------------
 @torch.inference_mode()
 def _ensure_qwen():
     global _q_proc, _q_model
     if _q_proc is None or _q_model is None:
-        _q_proc  = AutoProcessor.from_pretrained(MODEL_QWEN2_VL, trust_remote_code=True)
-        _q_model = AutoModelForVision2Seq.from_pretrained(MODEL_QWEN2_VL, torch_dtype=torch.float32, trust_remote_code=True)
+        _q_proc = AutoProcessor.from_pretrained(MODEL_QWEN2_VL, trust_remote_code=True)
+        _q_model = AutoModelForVision2Seq.from_pretrained(
+            MODEL_QWEN2_VL, torch_dtype=torch.float32, trust_remote_code=True
+        )
         _q_model.to(device()).eval()
 
+
 @torch.inference_mode()
-def caption_qwen(img: Image.Image, max_new_tokens: int = 220, temperature: float = 0.2, top_p: float = 0.9) -> str:
+def caption_qwen(img: Image.Image, max_new_tokens: int = 220, temperature: float = 0.3, top_p: float = 0.9) -> str:
     _ensure_qwen()
-    # Strong, explicit instruction for LONG descriptions
-    system_inst = (
-        "Describe this image in exhaustive detail. Include: subject(s), appearance, clothing, pose, "
-        "body details, interactions, accessories, objects, environment, background, textures, colors, "
-        "time of day, mood, and context. Be explicit, precise, and visual. Minimum ~80 words."
+    # Force long, explicit visual descriptions
+    inst = (
+        "Describe this image in exhaustive visual detail. Include: subjects, appearance, clothing, pose, "
+        "body details, accessories, objects, background, textures, colors, time of day, mood, context. "
+        "Be concrete and specific. Minimum ~80 words."
     )
-    inputs = _q_proc(text=system_inst, images=img, return_tensors="pt")
+    inputs = _q_proc(text=inst, images=img, return_tensors="pt")
     out = _q_model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
         do_sample=True,
         temperature=float(temperature),
         top_p=float(top_p),
-        repetition_penalty=1.05,
+        repetition_penalty=1.06,
     )
     return _q_proc.batch_decode(out, skip_special_tokens=True)[0].strip()
 
+
 @torch.inference_mode()
 def caption_qwen_retry_if_short(img: Image.Image, min_words: int, **gen_kwargs) -> str:
-    """Retry once with an even stronger instruction if the result is under min_words."""
-    text = caption_qwen(img, **gen_kwargs)
-    if len(text.split()) >= min_words:
-        return text
-    # Second pass: force structure and length
+    """Retry once with a stricter structured instruction if too short."""
+    first = caption_qwen(img, **gen_kwargs)
+    if len(first.split()) >= min_words:
+        return first
+
     _ensure_qwen()
     retry_inst = (
-        "Describe the image again with strict structure and more detail. Sections: "
-        "1) Subjects & actions, 2) Appearance & clothing, 3) Scene & background, "
-        "4) Lighting & colors, 5) Camera & framing, 6) Context & mood. "
-        "Write 120–180 words. Be concrete and visual."
+        "Rewrite a much longer, structured description of the image. Sections: "
+        "1) Subjects & actions; 2) Appearance & clothing; 3) Scene & background; "
+        "4) Lighting & colors; 5) Context & mood. Use 120–180 words. Be explicit and visual."
     )
     inputs = _q_proc(text=retry_inst, images=img, return_tensors="pt")
     out = _q_model.generate(
         **inputs,
         max_new_tokens=max(220, gen_kwargs.get("max_new_tokens", 220)),
         do_sample=True,
-        temperature=gen_kwargs.get("temperature", 0.2),
+        temperature=gen_kwargs.get("temperature", 0.3),
         top_p=gen_kwargs.get("top_p", 0.9),
-        repetition_penalty=1.07,
+        repetition_penalty=1.08,
     )
-    text2 = _q_proc.batch_decode(out, skip_special_tokens=True)[0].strip()
-    if len(text2.split()) > len(text.split()):
-        return text2
-    return text  # fallback to first if somehow shorter
+    second = _q_proc.batch_decode(out, skip_special_tokens=True)[0].strip()
+    return second if len(second.split()) >= min_words else first
+
 
 @torch.inference_mode()
 def _ensure_blip():
     global _b_proc, _b_model
     if _b_proc is None or _b_model is None:
-        _b_proc  = BlipProcessor.from_pretrained(MODEL_BLIP)
+        _b_proc = BlipProcessor.from_pretrained(MODEL_BLIP)
         _b_model = BlipForConditionalGeneration.from_pretrained(MODEL_BLIP)
         _b_model.to(device()).eval()
 
+
 @torch.inference_mode()
-def caption_blip(img: Image.Image, beam_search: bool = True, max_len: int = 96) -> str:
+def caption_blip(img: Image.Image, beam_search: bool = True, max_len: int = 140) -> str:
     _ensure_blip()
     inputs = _b_proc(images=img, return_tensors="pt")
-    max_len = int(max(32, min(160, max_len)))
+    max_len = int(max(32, min(200, max_len)))
     kwargs = dict(max_length=max_len, num_beams=4) if beam_search else dict(max_length=max_len)
     out = _b_model.generate(**inputs, **kwargs)
     return _b_proc.decode(out[0], skip_special_tokens=True).strip()
 
-def caption_image(model_choice: str, img: Image.Image, min_words: int, max_new_tokens: int):
-    if model_choice == "Qwen2-VL-2B-Instruct":
-        return caption_qwen_retry_if_short(
-            img,
-            min_words=min_words,
-            max_new_tokens=max_new_tokens,
-            temperature=0.35,
-            top_p=0.92,
-        )
-    else:
-        # BLIP won't reliably hit min_words; we just raise max_len and num_beams
-        base = caption_blip(img, beam_search=True, max_len=min(160, max_new_tokens))
-        return base
 
-# ---------------- Main generate ----------------
+def caption_image(model_choice: str, img: Image.Image, min_words: int, max_new_tokens: int) -> str:
+    # Guard invalid values coming from the UI
+    if model_choice not in VALID_MODELS:
+        model_choice = "Qwen2-VL-2B-Instruct"
+
+    # Try requested model, fall back to BLIP on any exception so UX doesn't hard-crash
+    try:
+        if model_choice == "Qwen2-VL-2B-Instruct":
+            return caption_qwen_retry_if_short(
+                img,
+                min_words=min_words,
+                max_new_tokens=max_new_tokens,
+                temperature=0.35,
+                top_p=0.92,
+            )
+        else:
+            # BLIP tends to be briefer; push max_len up
+            return caption_blip(img, beam_search=True, max_len=min(200, max_new_tokens))
+    except Exception:
+        # Fallback
+        try:
+            return caption_blip(img, beam_search=True, max_len=min(200, max_new_tokens))
+        except Exception as e:
+            # Last-resort: return readable error text in the caption box
+            return f"[Error] Captioning failed: {type(e).__name__}: {e}"
+
+
+# -------------------- Main generate --------------------
 def generate(
     image: Optional[Image.Image],
     text_scene: str,
@@ -195,10 +263,9 @@ def generate(
     max_new_tokens: int,
     seconds: int,
     seed: int,
-    model_choice: str
+    model_choice: str,
 ):
-    rng_seed = int(seed)
-    rnd = random.Random(rng_seed)
+    rnd = random.Random(int(seed))
 
     if image is not None:
         img = safe_open_image(image)
@@ -206,12 +273,14 @@ def generate(
         w, h = img.size
         exif = exif_info(img)
     else:
+        # Text-only path: use user text as the scene description directly
         if not text_scene or not text_scene.strip():
+            # Return empty strings so Gradio boxes clear but no crash
             return "", "", ""
         caption = text_scene.strip()
         w, h, exif = 1920, 1080, {}
 
-    angle, shot, lens, light, comp, style, post, neg, aspect = smart_choices(caption, w, h, exif, rng_seed)
+    angle, shot, lens, light, comp, style, post, neg, aspect = smart_choices(caption, w, h, exif, seed)
     image_prompt = build_image_prompt(caption, angle, shot, lens, light, comp, style, post, aspect)
     negative_line = f"\nNegative prompt hints: {neg}"
 
@@ -219,8 +288,8 @@ def generate(
         image_prompt,
         motion=rnd.choice(VIDEO_MOTION),
         meta=rnd.sample(VIDEO_META, k=2),
-        transition=rnd.choice(VIDEO_TRANS),
-        seconds=int(max(2, min(20, seconds)))
+        transition=rnd.choice(VIDEO_TRANSITIONS),
+        seconds=int(max(2, min(20, seconds))),
     )
 
     if mode == "Image prompt":
@@ -230,27 +299,29 @@ def generate(
     else:
         return caption, image_prompt + negative_line, video_prompt
 
-# ---------------- UI ----------------
+
+# -------------------- UI --------------------
 TITLE = "VisPrompt Studio — Long-Form Scene Descriptions"
 DESC = """
 Upload an image **or** describe a scene.  
-This version enforces **long, detailed scene descriptions** before building the camera/lighting prompts.
+This build forces **long, explicit scene descriptions** first, then adds detailed camera/lighting/comp info for image and video prompts.
 
-Models:
-- **Qwen2-VL-2B-Instruct (default)** — richer people/adult descriptions on CPU (slower).
-- **BLIP base** — faster but briefer.
+**Models**
+- Qwen2-VL-2B-Instruct (default): richer descriptions on people/adult photos (slower on CPU).
+- BLIP base: faster on CPU; simpler captions.
 """
 
 with gr.Blocks(title=TITLE) as demo:
     gr.Markdown(f"# {TITLE}\n{DESC}")
+
     with gr.Row():
         with gr.Column(scale=1):
-            model_choice = gr.Dropdown(
+            model_choice = gr.Radio(
                 choices=["Qwen2-VL-2B-Instruct", "BLIP base"],
                 value="Qwen2-VL-2B-Instruct",
-                label="Captioning model (drives the base scene description)",
-                allow_custom_value=False
+                label="Captioning model (drives the base scene description)"
             )
+
             with gr.Tab("Image input"):
                 image = gr.Image(type="pil", label="Upload an image", height=360)
             with gr.Tab("Text input"):
@@ -259,24 +330,39 @@ with gr.Blocks(title=TITLE) as demo:
             mode = gr.Radio(choices=["Image prompt", "Video prompt", "Both"], value="Both", label="Output type")
 
             with gr.Accordion("Description depth", open=True):
-                min_words = gr.Slider(40, 180, value=90, step=10, label="Min words for scene description (Qwen)")
-                max_new_tokens = gr.Slider(96, 300, value=220, step=4, label="Max tokens for caption generation")
+                min_words = gr.Slider(40, 200, value=90, step=10, label="Minimum words for scene description (Qwen)")
+                max_new_tokens = gr.Slider(96, 320, value=220, step=4, label="Max tokens for caption generation")
 
             with gr.Accordion("Video settings", open=False):
-                seconds = gr.Slider(3, 12, value=6, step=1, label="Video duration (seconds)")
+                seconds = gr.Slider(3, 12, value=6, step=1, label="Video duration for prompt (seconds)")
 
             seed = gr.Slider(0, 9999, value=1234, step=1, label="Style seed")
             run = gr.Button("Generate prompts", variant="primary")
 
         with gr.Column(scale=1):
-            caption_box = gr.Textbox(label="Long scene description (from image or your text)", interactive=False, lines=10, show_copy_button=True)
-            img_prompt_box = gr.Textbox(label="Image Prompt (copy to image generators)", interactive=False, lines=14, show_copy_button=True)
-            vid_prompt_box = gr.Textbox(label="Video Prompt (copy to video generators)", interactive=False, lines=14, show_copy_button=True)
+            caption_box = gr.Textbox(
+                label="Long scene description (from image or your text)",
+                interactive=False,
+                lines=10,
+                show_copy_button=True,
+            )
+            img_prompt_box = gr.Textbox(
+                label="Image Prompt (copy to image generators)",
+                interactive=False,
+                lines=14,
+                show_copy_button=True,
+            )
+            vid_prompt_box = gr.Textbox(
+                label="Video Prompt (copy to video generators)",
+                interactive=False,
+                lines=14,
+                show_copy_button=True,
+            )
 
     run.click(
         fn=generate,
         inputs=[image, text_scene, mode, min_words, max_new_tokens, seconds, seed, model_choice],
-        outputs=[caption_box, img_prompt_box, vid_prompt_box]
+        outputs=[caption_box, img_prompt_box, vid_prompt_box],
     )
 
 if __name__ == "__main__":