Upload app.py with huggingface_hub

app.py

@@ -34,7 +34,7 @@ from pathlib import Path
 import gradio as gr
 import torch
 from PIL import Image
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, InferenceClient
 
 from diffusers import ModularPipeline
 
@@ -48,6 +48,8 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 # when switching to the public klein repos.
 LINK_DISTILLED = os.environ.get("F3_LINK_DISTILLED", "https://huggingface.co/black-forest-labs/flux-3-image-klein")
 LINK_BASE = os.environ.get("F3_LINK_BASE", "https://huggingface.co/black-forest-labs/flux-3-image-klein-base")
+LINK_CODE = os.environ.get("F3_LINK_CODE", "https://github.com/black-forest-labs/flux-3-image-klein")
+LINK_BOXES = os.environ.get("F3_LINK_BOXES", "https://docs.bfl.ai/flux_3/flux3_image_bounding_boxes")
 
 MAX_SEED = 2**63 - 1
 DEFAULT_STEPS = int(os.environ.get("F3_STEPS", "16"))  # distilled default
@@ -169,6 +171,59 @@ def _decode_ref(data_uri):
     return img
 
 
+# ---- optional VLM prompt upsampling (klein-style) ---------------------------
+VLM_MODEL = os.environ.get("F3_VLM_MODEL", "Qwen/Qwen3-VL-235B-A22B-Instruct")
+_hf_client = InferenceClient(api_key=os.environ.get("INFERENCE_TOKEN") or HF_TOKEN)
+
+SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.3 by Black Forest Labs. Rewrite user prompts into a full scene caption while strictly preserving their core subject and intent.
+
+Guidelines:
+1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
+2. Details: Describe the subject, composition, lighting (quality, direction, color), atmosphere, textures, materials, spatial relationships, and environmental context.
+3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.
+
+Output only the revised prompt and nothing else."""
+
+SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.3 by Black Forest Labs, an image-editing expert. You turn an editing request into a complete caption of the target image: describe the reference image(s) in full with the requested edit folded in (60-120 words).
+
+Rules:
+- One caption only, no commentary
+- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
+- Keep everything the request does not change (faces, lighting, composition) and describe it as it is
+- Refer to the references as <ref_image_0>, <ref_image_1>, ... in the order given ("the person from <ref_image_0>")
+- Make abstractions concrete ("futuristic" -> "glowing cyan neon, metallic panels")
+- Keep content PG-13
+
+Output only the final caption in plain text and nothing else."""
+
+
+def _img_data_uri(img):
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+def _upsample_prompt(prompt, image_list=None):
+    """Rewrite the prompt into a full scene caption via a VLM; returns the original
+    prompt on any failure so generation always proceeds."""
+    try:
+        if image_list:
+            content = [{"type": "text", "text": prompt}]
+            for img in image_list:
+                content.append({"type": "image_url", "image_url": {"url": _img_data_uri(img)}})
+            messages = [{"role": "system", "content": SYSTEM_PROMPT_WITH_IMAGES},
+                        {"role": "user", "content": content}]
+        else:
+            messages = [{"role": "system", "content": SYSTEM_PROMPT_TEXT_ONLY},
+                        {"role": "user", "content": prompt}]
+        r = _hf_client.chat.completions.create(model=VLM_MODEL, messages=messages, max_tokens=1024)
+        out = (r.choices[0].message.content or "").strip()
+        return out or prompt
+    except Exception as e:
+        print(f"[app] prompt upsampling failed, using raw prompt: {e!r}", flush=True)
+        return prompt
+
+
 def _duration(state_json, *_):
     try:
         st = json.loads(state_json) if state_json else {}
@@ -191,7 +246,8 @@ def _generate(state_json):
     scene = st.get("scene") or ""
     aspect = st.get("aspect") or "1:1"
 
-    full_prompt = _build_prompt(scene, elements)
+    # `generate` may have pre-upsampled the prompt (off-GPU); use it if present.
+    full_prompt = st.get("_final_prompt") or _build_prompt(scene, elements)
     if not full_prompt.strip():
         raise gr.Error("Add at least a scene prompt or one described element.")
 
@@ -251,9 +307,15 @@ os.makedirs(CACHE_DIR, exist_ok=True)
 _CACHE_MEM = {}
 
 
+def _has_boxes(st):
+    return any((e.get("desc") or "").strip() for e in (st.get("elements") or []))
+
+
 def _cache_key(st):
     if st.get("ref_image"):
         return None  # i2i is not a fixed example
+    if st.get("prompt_upsampling") and not _has_boxes(st):
+        return None  # VLM upsampling is non-deterministic -> don't cache
     seed = st.get("seed")
     if not (isinstance(seed, (int, float)) and seed >= 0):
         return None  # random seed -> not a reproducible example
@@ -293,6 +355,17 @@ def generate(state_json):
         st = json.loads(state_json) if state_json else {}
     except Exception:
         st = {}
+    # Optional VLM prompt upsampling, done OFF-GPU before allocation. A prompt that
+    # carries a box table (layout) is sent as-is, matching the klein space.
+    if st.get("prompt_upsampling") and not _has_boxes(st):
+        base = _build_prompt(st.get("scene") or "", st.get("elements") or [])
+        if base.strip():
+            ref = _decode_ref(st.get("ref_image"))
+            up = _upsample_prompt(base, [ref] if ref is not None else None)
+            if up and up.strip():
+                st["_final_prompt"] = up
+                state_json = json.dumps(st)
+                print(f"[app] upsampled prompt -> {up[:90]!r}", flush=True)
     key = _cache_key(st)
     if key:
         cached = _cache_get(key)
@@ -380,7 +453,8 @@ RATIO_LIST = list(RATIOS.keys())
 CONFIG_JS = (
     "<script>window.F3_PRESETS = " + json.dumps(PRESETS) + ";"
     "window.F3_RATIOS = " + json.dumps(RATIO_LIST) + ";"
-    "window.F3_LINKS = " + json.dumps({"distilled": LINK_DISTILLED, "base": LINK_BASE}) + ";"
+    "window.F3_LINKS = " + json.dumps({"distilled": LINK_DISTILLED, "base": LINK_BASE,
+                                       "code": LINK_CODE, "boxes": LINK_BOXES}) + ";"
     "window.F3_PRESET_ORDER = " + json.dumps(list(PRESETS.keys())) + ";</script>"
 )
 
@@ -510,12 +584,18 @@ STYLE2 = """
 #f3-root .titlebar .tb-text{min-width:0;}
 #f3-root .titlebar h1{margin:0 0 4px;font-size:22px;font-weight:650;color:var(--txt);letter-spacing:-.01em;}
 #f3-root .titlebar p{margin:0;color:var(--dim);font-size:13.5px;max-width:760px;line-height:1.5;}
-/* klein-style pill links */
+/* klein-style button links */
 #f3-root .f3-links{display:flex;gap:8px;flex-wrap:wrap;}
-#f3-root .f3-links a,#f3-root .f3-links a:visited{text-decoration:none;font-size:12.5px;padding:6px 13px;border-radius:999px;
+#f3-root .f3-links a,#f3-root .f3-links a:visited{text-decoration:none;font-size:12.5px;padding:6px 12px;border-radius:8px;
   border:1px solid var(--line2);color:var(--txt);background:var(--panel);white-space:nowrap;transition:border-color .12s,color .12s;}
 #f3-root .f3-links a:hover{border-color:var(--violet);color:var(--violet);}
 
+/* prompt-upsampling toggle (under the prompt in the input bar) */
+#f3-root .ib-up{margin-top:9px;}
+#f3-root .ib-up .btn{font-size:12px;color:var(--dim);}
+#f3-root .ib-up .btn.on{color:var(--txt);}
+#f3-root .ib-up .btn .hintq{color:var(--dim2);font-size:11px;}
+
 /* reference image block (i2i) */
 #f3-root .refblock{padding:14px 18px;border-bottom:1px solid var(--line2);}
 #f3-root .refblock>label{display:block;font:600 11px/1.4 var(--mono);letter-spacing:.06em;color:var(--dim2);margin-bottom:9px;}
@@ -583,7 +663,7 @@ SCRIPT = r"""
     ["Prompt","One line describing the whole scene — or the edit, if you added a reference image."],
   ];
   let S = {aspect:"1:1", elements:[], scene:"", seed:-1, rendered:false, step:1, showBoxes:true,
-           refImage:null, refW:0, refH:0};
+           refImage:null, refW:0, refH:0, upsample:false};
   let seq = 0, sel = null;
 
   function clamp(v,a,b){return Math.max(a,Math.min(b,v));}
@@ -608,7 +688,7 @@ SCRIPT = r"""
   }
   function syncState(){
     setHidden('f3-state', JSON.stringify({
-      aspect:S.aspect, scene:S.scene, seed:S.seed, ref_image:S.refImage,
+      aspect:S.aspect, scene:S.scene, seed:S.seed, ref_image:S.refImage, prompt_upsampling:S.upsample,
       elements:S.elements.map(e=>({id:e.id,bbox:e.bbox,desc:e.desc,color:e.color})),
     }));
   }
@@ -899,8 +979,10 @@ SCRIPT = r"""
     +   '<p>Draw a box for every element, describe it, and the step-distilled FLUX.3 model renders your layout. '
     +   'Upload a reference image to edit it with the same boxes.</p></div>'
     + '<div class="f3-links">'
-    +   '<a href="'+((window.F3_LINKS||{}).distilled||'#')+'" target="_blank" rel="noopener">Distilled model ↗</a>'
-    +   '<a href="'+((window.F3_LINKS||{}).base||'#')+'" target="_blank" rel="noopener">Base model ↗</a>'
+    +   '<a href="'+((window.F3_LINKS||{}).distilled||'#')+'" target="_blank" rel="noopener">Distilled weights ↗</a>'
+    +   '<a href="'+((window.F3_LINKS||{}).base||'#')+'" target="_blank" rel="noopener">Base weights ↗</a>'
+    +   '<a href="'+((window.F3_LINKS||{}).code||'#')+'" target="_blank" rel="noopener">Code ↗</a>'
+    +   '<a href="'+((window.F3_LINKS||{}).boxes||'#')+'" target="_blank" rel="noopener">Bounding boxes ↗</a>'
     + '</div>'
     +'</div>'
     // INPUT BAR — the type-here / drop-here controls, in their own distinct zone above the canvas
@@ -912,7 +994,9 @@ SCRIPT = r"""
     +       '<button class="refclear">Remove</button></div></div>'
     +     '<input type="file" accept="image/*" class="reffile" style="display:none"></div>'
     +   '<div class="ib-prompt"><label class="ib-lab">PROMPT</label>'
-    +     '<textarea rows="2" placeholder="Describe the whole scene in one line — or the edit, if you added a reference image."></textarea></div>'
+    +     '<textarea rows="2" placeholder="Describe the whole scene in one line — or the edit, if you added a reference image."></textarea>'
+    +     '<div class="ib-up"><span class="btn upBtn"><span class="toggle"></span>Enhance prompt '
+    +       '<span class="hintq">· VLM, text prompts only</span></span></div></div>'
     +   '<div class="ib-gen"><label class="ib-lab" style="visibility:hidden">GO</label><button class="genbtn">Generate</button></div>'
     +'</div>'
     +'<div class="main">'
@@ -964,6 +1048,7 @@ SCRIPT = r"""
       replayBtn:root.querySelector('.replayBtn'), cnt:root.querySelector('.table-bar .cnt'),
       etable:root.querySelector('.etable'),
       reuseBtn:root.querySelector('.reuseBtn'), clearBtn:root.querySelector('.clearBtn'),
+      upBtn:root.querySelector('.upBtn'),
       refBox:root.querySelector('.refbox'), refThumb:root.querySelector('.refthumb'),
       reffile:root.querySelector('.reffile') };
 
@@ -974,6 +1059,8 @@ SCRIPT = r"""
     NS.el.clearBtn.addEventListener('click',clearCanvas);
     NS.el.reuseBtn.addEventListener('click',reuseAsInput);
     NS.el.toggleBtn.addEventListener('click',()=>{ S.showBoxes=!S.showBoxes; updateToggle(); });
+    NS.el.upBtn.addEventListener('click',()=>{ S.upsample=!S.upsample;
+      NS.el.upBtn.classList.toggle('on',S.upsample); syncState(); });
     root.querySelector('.table-bar').addEventListener('click',()=>NS.el.etable.classList.toggle('open'));
     window.addEventListener('resize',()=>{ layoutBoard(); });