Align Prompt Enhancer with try-on safetensors inference

README.md

@@ -61,9 +61,23 @@ JOY_TRANSFORMER_PATH=transformer
 JOY_DIT_CKPT_TYPE=safetensor
 JOY_MODEL_PATH=FireRed-Image-Edit-1.0
 JOY_VAE_PATH=FireRed-Image-Edit-1.0/vae/diffusion_pytorch_model.safetensors
+QWEN235_CHAT_URL=http://ai-api.jdcloud.com/v1/chat/completions
+QWEN235_MODEL=qwen3-vl-235
 HF_HOME=/data/.huggingface
 ```
 
+For Prompt Enhancer, add one of these as a Space Secret:
+
+```text
+QWEN235_API_KEY=your_api_key
+```
+
+or:
+
+```text
+OPENAI_API_KEY=your_api_key
+```
+
 ## Expected Model Repository Layout
 
 This wrapper downloads `MODEL_ID` with `huggingface_hub.snapshot_download`.
@@ -108,6 +122,8 @@ prompt = <image><image> + try-on instruction
 output = one generated PIL image
 ```
 
+Prompt Enhancer is enabled by default in the UI. It sends the garment image and person image to the configured Qwen3-VL-235 endpoint, then displays the enhanced prompt used for generation.
+
 If multiple garment images are uploaded, the app uses `primary_garment_index` to select one primary garment. It does not pretend to perform multi-reference fusion.
 
 ## Deployment Steps

app.py

@@ -29,6 +29,7 @@ def predict(
     garment_type: str,
     primary_garment_index: int,
     prompt: str,
+    prompt_enhancer: bool,
     negative_prompt: str,
     seed: int,
     randomize_seed: bool,
@@ -45,7 +46,7 @@ def predict(
         person = resize_for_demo(person_image, max_side=1024)
         garments = [resize_for_demo(img, max_side=1024) for img in garment_images]
 
-        result, used_index, used_prompt, status = run_tryon(
+        result, used_index, enhanced_prompt, pe_status, status = run_tryon(
             person_image=person,
             garment_images=garments,
             garment_type=garment_type,
@@ -55,10 +56,11 @@ def predict(
             guidance_scale=float(guidance_scale),
             auto_crop=bool(crop_and_paste),
             prompt=prompt,
+            prompt_enhancer=bool(prompt_enhancer),
             negative_prompt=negative_prompt,
         )
 
-        return result, int(seed), used_index, used_prompt, status
+        return result, int(seed), used_index, enhanced_prompt, pe_status, status
 
     except gr.Error:
         raise
@@ -108,11 +110,12 @@ with gr.Blocks(title="Virtual Try-On ZeroGPU Demo") as demo:
             )
 
             prompt = gr.Textbox(
-                label="Try-On Prompt",
+                label="Original Try-On Prompt",
                 value="",
                 lines=3,
                 placeholder="Leave empty to use the default prompt for the selected garment type.",
             )
+            prompt_enhancer = gr.Checkbox(value=True, label="Prompt Enhancer")
 
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Number(value=42, precision=0, label="Seed")
@@ -148,7 +151,8 @@ with gr.Blocks(title="Virtual Try-On ZeroGPU Demo") as demo:
             output_image = gr.Image(label="Try-On Result", type="pil")
             used_seed = gr.Number(label="Used Seed", precision=0)
             used_garment_index = gr.Number(label="Used Garment Index", precision=0)
-            used_prompt = gr.Textbox(label="Prompt Used", interactive=False, lines=3)
+            enhanced_prompt = gr.Textbox(label="Enhanced Prompt Used", interactive=False, lines=5)
+            pe_status = gr.Textbox(label="Prompt Enhancer Status", interactive=False)
             status = gr.Textbox(label="Status", interactive=False)
 
     run_button.click(
@@ -159,6 +163,7 @@ with gr.Blocks(title="Virtual Try-On ZeroGPU Demo") as demo:
             garment_type,
             primary_garment_index,
             prompt,
+            prompt_enhancer,
             negative_prompt,
             seed,
             randomize_seed,
@@ -166,7 +171,7 @@ with gr.Blocks(title="Virtual Try-On ZeroGPU Demo") as demo:
             guidance_scale,
             crop_and_paste,
         ],
-        outputs=[output_image, used_seed, used_garment_index, used_prompt, status],
+        outputs=[output_image, used_seed, used_garment_index, enhanced_prompt, pe_status, status],
         api_name="tryon",
     )
 

inference.py

@@ -1,15 +1,22 @@
 from __future__ import annotations
 
+import base64
+import hashlib
+import io
+import json
 import math
 import os
 import random
+import re
 import sys
 import contextlib
+import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import numpy as np
+import requests
 import torch
 import torchvision.transforms.functional as TF
 from huggingface_hub import snapshot_download
@@ -32,6 +39,28 @@ DEFAULT_PROMPTS = {
     "full_body": "Make the person in Picture 2 wear the clothing from Picture 1. Keep the pose, body shape, background, and lighting from Picture 2.",
 }
 
+QWEN235_MODEL = os.getenv("QWEN235_MODEL", "qwen3-vl-235")
+QWEN235_CHAT_URL = os.getenv("QWEN235_CHAT_URL", "http://ai-api.jdcloud.com/v1/chat/completions")
+PE_SYSTEM_ROLE = "You are an expert vision AI prompt engineer."
+PE_TEMPLATE_DEFAULT = (
+    "# Multi-Image Edit Instruction Rewriter\n\n"
+    "You are an expert at refining MULTI-IMAGE editing instructions. Given the "
+    "user's original prompt and the uploaded reference images (in order: Image 1, "
+    "Image 2, Image 3, Image 4), produce a single improved English editing "
+    "caption that a diffusion-based multi-image editing model can follow precisely.\n\n"
+    "## Core Rules\n"
+    "1. Preserve intent. Never invent new operations. Only clarify, enrich and correct what the user already asked for.\n"
+    "2. Always refer to the inputs as Image 1 / Image 2 / Image 3 / Image 4.\n"
+    "3. For every visual element in the output, state explicitly which reference image it comes from.\n"
+    "4. Output the rewritten English caption only. No preamble, no JSON, no markdown fences, no explanations.\n\n"
+    "## Replacement / Virtual Try-On Rules\n"
+    "For try-on, consistency is more important than creativity. Do not add cinematic effects, extra props, "
+    "new lighting, new camera angles, or creative embellishments. Preserve everything in the person image "
+    "that is not explicitly replaced: identity, face, hair, skin tone, body shape, pose, expression, "
+    "background, lighting, color tone, camera angle, and composition. Replace only the garment, preserving "
+    "the garment reference's color, pattern, texture, logo, silhouette, and natural fabric drape."
+)
+
 
 @dataclass
 class LoadedTryOnModel:
@@ -43,6 +72,7 @@ class LoadedTryOnModel:
 
 _MODEL: LoadedTryOnModel | None = None
 _MODEL_DIR: Path | None = None
+_PE_CACHE: dict[str, str] = {}
 
 
 def _get_bool_env(name: str, default: bool = False) -> bool:
@@ -70,6 +100,155 @@ def set_seed(seed: int) -> torch.Generator:
     return torch.Generator(device="cuda").manual_seed(seed)
 
 
+def _encode_image_base64_png(pil_image: Image.Image) -> str:
+    buf = io.BytesIO()
+    pil_image.convert("RGB").save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+
+def _remove_think_block(text: str) -> str:
+    return re.sub(r"<think>.*?</think>", "", text or "", flags=re.IGNORECASE | re.DOTALL).strip()
+
+
+def _normalize_chat_message_content(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content.strip()
+
+    texts: list[str] = []
+
+    def walk(value: Any) -> None:
+        if value is None:
+            return
+        if isinstance(value, str):
+            if value.strip():
+                texts.append(value.strip())
+            return
+        if isinstance(value, dict):
+            text = value.get("text")
+            if isinstance(text, str) and text.strip():
+                texts.append(text.strip())
+            for child in value.values():
+                if isinstance(child, (dict, list)):
+                    walk(child)
+            return
+        if isinstance(value, list):
+            for item in value:
+                walk(item)
+
+    walk(content)
+    return "\n".join(texts).strip()
+
+
+def _extract_caption_text(text: str) -> str:
+    clean = _remove_think_block((text or "").strip())
+    if not clean:
+        return ""
+    clean = clean.replace("```json", "").replace("```", "").strip()
+    try:
+        obj = json.loads(clean)
+        if isinstance(obj, dict) and isinstance(obj.get("Rewritten"), str):
+            return obj["Rewritten"].strip().replace("\n", " ")
+    except Exception:
+        pass
+    return clean
+
+
+def _image_signature(image: Image.Image | None) -> str:
+    if image is None:
+        return "none"
+    try:
+        buf = io.BytesIO()
+        image.convert("RGB").save(buf, format="PNG")
+        data = buf.getvalue()
+        digest = hashlib.md5(data[:65536]).hexdigest()
+        return f"{image.size[0]}x{image.size[1]}_{digest}"
+    except Exception:
+        return f"{image.size[0]}x{image.size[1]}"
+
+
+def _build_pe_prompt_text(original_prompt: str) -> str:
+    original_prompt = str(original_prompt or "").strip()
+    if not original_prompt:
+        return PE_TEMPLATE_DEFAULT
+    return (
+        f"Based on the uploaded image(s) and the prompt: '{original_prompt}', "
+        "please follow the instructions below to produce the rewritten caption.\n\n"
+        f"{PE_TEMPLATE_DEFAULT}"
+    )
+
+
+def _rewrite_prompt_qwen3vl235_chatapi(prompt: str, images: list[Image.Image]) -> str:
+    prompt = str(prompt or "").strip()
+    if not prompt:
+        return prompt
+
+    api_key = os.getenv("QWEN235_API_KEY") or os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise EnvironmentError("QWEN235_API_KEY or OPENAI_API_KEY is not set")
+
+    user_content: list[dict[str, Any]] = []
+    for image in images:
+        user_content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/png;base64,{_encode_image_base64_png(image)}"},
+            }
+        )
+    user_content.append({"type": "text", "text": _build_pe_prompt_text(prompt)})
+
+    payload = {
+        "model": QWEN235_MODEL,
+        "messages": [
+            {"role": "system", "content": PE_SYSTEM_ROLE},
+            {"role": "user", "content": user_content},
+        ],
+        "max_tokens": int(os.getenv("QWEN235_MAX_TOKENS", "1024")),
+        "stream": False,
+    }
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+        "Trace-Id": f"tryon-space-qwen235-{int(time.time() * 1000)}",
+    }
+
+    last_err: Exception | None = None
+    for attempt in range(int(os.getenv("QWEN235_MAX_RETRIES", "3"))):
+        try:
+            resp = requests.post(QWEN235_CHAT_URL, headers=headers, json=payload, timeout=float(os.getenv("QWEN235_TIMEOUT", "60")))
+            if resp.status_code != 200:
+                raise RuntimeError(f"Qwen3-VL-235 API failed: status={resp.status_code}, text={resp.text[:2000]}")
+            data = resp.json()
+            content = (((data.get("choices") or [{}])[0].get("message") or {}).get("content"))
+            caption = _extract_caption_text(_normalize_chat_message_content(content))
+            return caption or prompt
+        except Exception as exc:
+            last_err = exc
+            time.sleep(0.5 * (2**attempt))
+
+    raise RuntimeError(f"Prompt Enhancer failed: {last_err}")
+
+
+def rewrite_prompt_if_enabled(prompt: str, images: list[Image.Image], enabled: bool) -> tuple[str, str]:
+    prompt = str(prompt or "")
+    if not enabled:
+        return prompt, "Prompt Enhancer disabled."
+    if not images:
+        return prompt, "Prompt Enhancer skipped: no reference images."
+
+    cache_key = f"qwen235|p={prompt.strip()}|imgs={'|'.join(_image_signature(img) for img in images)}"
+    if cache_key in _PE_CACHE:
+        return _PE_CACHE[cache_key], "Prompt Enhancer cache hit."
+
+    try:
+        rewritten = _rewrite_prompt_qwen3vl235_chatapi(prompt, [img.convert("RGB") for img in images])
+        _PE_CACHE[cache_key] = rewritten
+        return rewritten, "Prompt Enhancer applied."
+    except Exception as exc:
+        return prompt, f"Prompt Enhancer failed, using original prompt: {exc}"
+
+
 def _download_model_repo() -> Path:
     model_id = os.getenv("MODEL_ID")
     if not model_id:
@@ -504,7 +683,8 @@ def run_tryon(
     auto_crop: bool = False,
     prompt: str | None = None,
     negative_prompt: str = DEFAULT_NEG_PROMPT,
-) -> tuple[Image.Image, int, str, str]:
+    prompt_enhancer: bool = True,
+) -> tuple[Image.Image, int, str, str, str]:
     if person_image is None:
         raise ValueError("person_image is required.")
     if not garment_images:
@@ -517,6 +697,7 @@ def run_tryon(
     person = person_image.convert("RGB")
     garment = primary_garment.convert("RGB")
     effective_prompt = _build_prompt(prompt, garment_type)
+    enhanced_prompt, pe_status = rewrite_prompt_if_enabled(effective_prompt, [garment, person], prompt_enhancer)
     effective_negative = (negative_prompt or DEFAULT_NEG_PROMPT).strip()
 
     if auto_crop:
@@ -524,7 +705,7 @@ def run_tryon(
             model,
             person,
             garment,
-            effective_prompt,
+            enhanced_prompt,
             num_inference_steps,
             guidance_scale,
             seed,
@@ -533,7 +714,7 @@ def run_tryon(
         if cropped_result is not None:
             image, path_status = cropped_result
             status = f"Success. Used garment index {used_index}. {path_status}."
-            return image, used_index, effective_prompt, status
+            return image, used_index, enhanced_prompt, pe_status, status
 
     target_h, target_w = get_bucket_size(person)
     input_person = resize_center_crop(person, (target_h, target_w))
@@ -541,7 +722,7 @@ def run_tryon(
     output = _run_pipeline(
         model,
         [input_garment, input_person],
-        effective_prompt,
+        enhanced_prompt,
         target_h,
         target_w,
         num_inference_steps,
@@ -550,7 +731,7 @@ def run_tryon(
         effective_negative,
     )
     status = f"Success. Used garment index {used_index}. direct path."
-    return output, used_index, effective_prompt, status
+    return output, used_index, enhanced_prompt, pe_status, status
 
 
 def clear_model_cache() -> None:

requirements.txt

@@ -9,6 +9,7 @@ transformers
 safetensors
 pillow
 opencv-python-headless
+requests
 numpy
 scipy
 einops