Update app.py

app.py

@@ -1,125 +1,85 @@
 # app.py
-# VisPrompt Studio — Image/Text → Detailed Prompts (CPU, uncensored)
-# - Runs on Hugging Face Spaces (free CPU)
-# - Upload an image OR type a scene; get an image prompt, a video prompt, or both
-# - Uses Salesforce/blip-image-captioning-base for image understanding (no safety filters)
+# VisPrompt Studio — Image/Text → Detailed Prompts (CPU, mostly uncensored)
+# Models:
+#  - Qwen2-VL-2B-Instruct  (default; stronger on people/NSFW, slower on CPU)
+#  - Salesforce/blip-image-captioning-base (fallback; faster on CPU)
 
-import os
-import random
-import math
-from typing import Tuple, Optional
+import os, random, math
+from typing import Optional, Tuple
 
 import gradio as gr
 from PIL import Image, ExifTags
 import torch
-from transformers import BlipProcessor, BlipForConditionalGeneration
 
-HF_MODEL_ID = "Salesforce/blip-image-captioning-base"
-
-# Lazy globals
-_processor: Optional[BlipProcessor] = None
-_model: Optional[BlipForConditionalGeneration] = None
-
-
-def _load_blip():
-    """Load BLIP processor/model into CPU once."""
-    global _processor, _model
-    if _processor is None or _model is None:
-        _processor = BlipProcessor.from_pretrained(HF_MODEL_ID)
-        _model = BlipForConditionalGeneration.from_pretrained(HF_MODEL_ID)
-        _model.to("cpu")
-        _model.eval()
-
-
-def safe_open_image(inp) -> Optional[Image.Image]:
-    if inp is None:
-        return None
-    if isinstance(inp, Image.Image):
-        return inp.convert("RGB")
+from transformers import (
+    AutoProcessor,
+    AutoModelForVision2Seq,
+    BlipProcessor,
+    BlipForConditionalGeneration,
+)
+
+# ===== Model IDs =====
+MODEL_QWEN2_VL = "Qwen/Qwen2-VL-2B-Instruct"
+MODEL_BLIP = "Salesforce/blip-image-captioning-base"
+
+# Lazy globals for each backend
+_q_processor = None
+_q_model = None
+_b_processor = None
+_b_model = None
+
+def device():
+    # Free CPU Space: force cpu. (If you later use ZeroGPU, you can set CUDA here.)
+    return "cpu"
+
+# ---------- Utilities ----------
+def safe_open_image(inp):
+    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:
-    """Extract a couple of useful EXIF fields if present."""
     out = {}
     try:
         raw = img.getexif()
-        if not raw:
-            return out
+        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)
+        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 "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
 
-
-# Small lexicons to enrich prompts (kept inline for mobile copy/paste simplicity)
-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"
-]
-
+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):
-    """Pick plausible settings using simple heuristics."""
     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"]):
+    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"]):
+    if any(k in text 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 text for k in ["towering","imposing","hero","statue"]):
         angle = "low-angle"
     else:
-        angle = rnd.choice(CAMERA_ANGLES[:3])  # eye/high/low
-
+        angle = rnd.choice(CAMERA_ANGLES[:3])
     lens = f'{exif["focal_length_mm"]}mm' if "focal_length_mm" in exif else rnd.choice(LENSES)
     light = rnd.choice(LIGHTING)
     comp = ", ".join(rnd.sample(COMPOSITION, k=2))
@@ -128,61 +88,60 @@ def smart_choices(caption: str, w: int, h: int, exif: dict, seed: int):
     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")
+    return (f"{desc}. Camera angle: {angle}. Shot size: {shot}. Lens: {lens}. "
+            f"Lighting: {light}. Composition: {comp}. Style: {style}. Post-processing: {post}. "
+            f"Aspect ratio hint: {ar}. Ultra-detailed, physically plausible, consistent scale.")
 
-@torch.inference_mode()
-def caption_image(img: Image.Image, beam_search: bool, max_len: int) -> str:
-    _load_blip()
-    inputs = _processor(images=img, return_tensors="pt")
-    max_len = int(max(16, min(128, max_len)))
-    if beam_search:
-        kwargs = dict(max_length=max_len, num_beams=3)
-    else:
-        kwargs = dict(max_length=max_len)
-    out = _model.generate(**inputs, **kwargs)
-    return _processor.decode(out[0], skip_special_tokens=True).strip()
-
-
-def build_image_prompt(subject_desc: str, angle: str, shot: str, lens: str,
-                       light: str, comp: str, style: str, post: str, aspect: float) -> str:
-    ar_text = "16:9" if aspect >= 1.6 else ("4:3" if aspect >= 1.3 else "1:1")
-    prompt = (
-        f"{subject_desc}. "
-        f"Camera angle: {angle}. Shot size: {shot}. Lens: {lens}. "
-        f"Lighting: {light}. Composition: {comp}. Style: {style}. Post-processing: {post}. "
-        f"Aspect ratio hint: {ar_text}. "
-        f"Ultra-detailed, physically plausible, consistent scale."
-    )
-    return prompt
+def build_video_prompt(img_prompt: str, motion: str, meta: list, transition: str, seconds: int):
+    meta_str = ", ".join(meta)
+    return (f"{img_prompt} Now render as a {seconds}s cinematic video; camera motion: {motion}; "
+            f"{meta_str}; transitions: {transition}. Maintain temporal consistency, realistic parallax, coherent lighting.")
 
+# ---------- Backends ----------
+@torch.inference_mode()
+def qwen2_vl_caption(img: Image.Image, max_new_tokens: int = 96) -> str:
+    global _q_processor, _q_model
+    if _q_processor is None or _q_model is None:
+        _q_processor = 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()
+    # Simple instruction to push detailed descriptions (uncensored)
+    prompt = "Describe this image in exhaustive detail: people, clothing, actions, camera angle, lighting, background."
+    inputs = _q_processor(text=prompt, images=img, return_tensors="pt")
+    out = _q_model.generate(**inputs, max_new_tokens=max_new_tokens)
+    text = _q_processor.batch_decode(out, skip_special_tokens=True)[0].strip()
+    return text
 
-def build_video_prompt(image_prompt: str, motion: str, meta: list, transition: str, seconds: int) -> str:
-    meta_str = ", ".join(meta)
-    return (
-        f"{image_prompt} "
-        f"Now render as a {seconds}s cinematic video; camera motion: {motion}; "
-        f"{meta_str}; transitions: {transition}. "
-        f"Maintain temporal consistency, realistic motion parallax, and coherent lighting through the sequence."
-    )
+@torch.inference_mode()
+def blip_caption(img: Image.Image, beam_search: bool = True, max_len: int = 80) -> str:
+    global _b_processor, _b_model
+    if _b_processor is None or _b_model is None:
+        _b_processor = BlipProcessor.from_pretrained(MODEL_BLIP)
+        _b_model = BlipForConditionalGeneration.from_pretrained(MODEL_BLIP)
+        _b_model.to(device()).eval()
+    inputs = _b_processor(images=img, return_tensors="pt")
+    max_len = int(max(16, min(128, max_len)))
+    kwargs = dict(max_length=max_len, num_beams=3) if beam_search else dict(max_length=max_len)
+    out = _b_model.generate(**inputs, **kwargs)
+    return _b_processor.decode(out[0], skip_special_tokens=True).strip()
 
+def caption_image(model_choice: str, img: Image.Image, beam_search: bool, max_len: int) -> str:
+    if model_choice == "Qwen2-VL-2B-Instruct":
+        return qwen2_vl_caption(img, max_new_tokens=max_len)
+    else:
+        return blip_caption(img, beam_search=beam_search, max_len=max_len)
 
-def generate(
-    image: Optional[Image.Image],
-    text_scene: str,
-    mode: str,
-    beam_search: bool,
-    max_len: int,
-    seconds: int,
-    seed: int
-):
-    """
-    Returns: (caption, image_prompt, video_prompt) per selected mode.
-    """
+# ---------- Main generate ----------
+def generate(image: Optional[Image.Image], text_scene: str, mode: str,
+             beam_search: bool, max_len: int, seconds: int, seed: int, model_choice: str):
     img = safe_open_image(image)
     caption = ""
     base_desc = ""
 
     if img is not None:
-        caption = caption_image(img, beam_search=beam_search, max_len=max_len)
+        caption = caption_image(model_choice, img, beam_search, max_len)
         w, h = img.size
         exif = exif_info(img)
         angle, shot, lens, light, comp, style, post, neg, aspect = smart_choices(caption, w, h, exif, seed)
@@ -196,8 +155,7 @@ def generate(
         angle, shot, lens, light, comp, style, post, neg, aspect = smart_choices(base_desc, w, h, exif, seed)
 
     image_prompt = build_image_prompt(base_desc, angle, shot, lens, light, comp, style, post, aspect)
-    negative_line = f"\nNegative prompt hints: {neg}"
-
+    neg_line = f"\nNegative prompt hints: {neg}"
     video_prompt = build_video_prompt(
         image_prompt,
         motion=random.choice(VIDEO_MOTION),
@@ -207,36 +165,36 @@ def generate(
     )
 
     if mode == "Image prompt":
-        return caption, image_prompt + negative_line, ""
+        return caption, image_prompt + neg_line, ""
     elif mode == "Video prompt":
         return caption, "", video_prompt
     else:
-        return caption, image_prompt + negative_line, video_prompt
+        return caption, image_prompt + neg_line, video_prompt
 
-
-TITLE = "VisPrompt Studio — Image/Text → Prompts (CPU)"
+# ---------- UI ----------
+TITLE = "VisPrompt Studio — Image/Text → Prompts (Qwen2-VL or BLIP)"
 DESC = """
 Upload an image **or** describe a scene.  
-Get a polished **image prompt**, a **video prompt**, or **both** — with camera angle, shot size, lens, lighting, composition, style, post, aspect ratio hints, plus optional negative prompt hints.
+Get a polished **image prompt**, a **video prompt**, or **both** — with camera angle, shot size, lens, lighting, composition, style, post, AR hints, plus optional negative prompts.
 
-- **Model:** Salesforce/blip-image-captioning-base (CPU; no built-in safety filters)
-- **Hardware:** Works on free CPU Spaces (first run downloads weights)
+**Models**:  
+- **Qwen2-VL-2B-Instruct (default)** — stronger on people/NSFW & fashion; 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(["Qwen2-VL-2B-Instruct","BLIP base"], value="Qwen2-VL-2B-Instruct", label="Captioning model")
             with gr.Tab("Image input"):
                 image = gr.Image(type="pil", label="Upload an image", height=360)
             with gr.Tab("Text input"):
                 text_scene = gr.Textbox(lines=6, placeholder="Describe the scene you want to reproduce...", label="Scene description")
 
-            mode = gr.Radio(choices=["Image prompt", "Video prompt", "Both"], value="Both", label="Output type")
-            beam_search = gr.Checkbox(True, label="Use beam search for caption (slightly slower, more detail)")
-            max_len = gr.Slider(48, 128, value=80, step=1, label="Caption max length")
+            mode = gr.Radio(choices=["Image prompt","Video prompt","Both"], value="Both", label="Output type")
+            beam_search = gr.Checkbox(True, label="BLIP: beam search (ignored for Qwen2-VL)")
+            max_len = gr.Slider(48, 128, value=96, step=1, label="Caption length / max tokens")
             seconds = gr.Slider(3, 12, value=6, step=1, label="Video duration (seconds)")
             seed = gr.Slider(0, 9999, value=1234, step=1, label="Style seed")
 
@@ -244,24 +202,14 @@ with gr.Blocks(title=TITLE) as demo:
 
         with gr.Column(scale=1):
             caption_box = gr.Textbox(label="Image caption (from model)", interactive=False, lines=3, show_copy_button=True)
-            img_prompt_box = gr.Textbox(
-                label="Image Prompt (copy-paste to image generators)",
-                interactive=False,
-                lines=12,
-                show_copy_button=True,
-            )
-            vid_prompt_box = gr.Textbox(
-                label="Video Prompt (copy-paste to video generators)",
-                interactive=False,
-                lines=14,
-                show_copy_button=True,
-            )
+            img_prompt_box = gr.Textbox(label="Image Prompt", interactive=False, lines=12, show_copy_button=True)
+            vid_prompt_box = gr.Textbox(label="Video Prompt", interactive=False, lines=14, show_copy_button=True)
 
     run.click(
         fn=generate,
-        inputs=[image, text_scene, mode, beam_search, max_len, seconds, seed],
+        inputs=[image, text_scene, mode, beam_search, max_len, seconds, seed, model_choice],
         outputs=[caption_box, img_prompt_box, vid_prompt_box]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()