| import gradio as gr |
| import numpy as np |
| import random |
| import torch |
| import spaces |
|
|
| from huggingface_hub import hf_hub_download |
| from accelerate import init_empty_weights |
| from accelerate.utils import set_module_tensor_to_device |
| from safetensors import safe_open |
|
|
| from PIL import Image |
| from torchao.quantization import Float8DynamicActivationFloat8WeightConfig |
| from torchao.quantization import quantize_ |
| from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline |
| from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel |
| from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3 |
|
|
| import time |
|
|
|
|
| def update_history(new_images, history): |
| time.sleep(0.5) |
| if history is None: |
| history = [] |
| if new_images is not None and len(new_images) > 0: |
| if not isinstance(history, list): |
| history = list(history) if history else [] |
| for img in new_images: |
| history.insert(0, img) |
| history = history[:20] |
| return history |
|
|
| def use_history_as_input(evt: gr.SelectData): |
| if evt.value is not None: |
| return gr.update(value=evt.value) |
| return gr.update() |
|
|
| dtype = torch.bfloat16 |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511" |
| BASE_MODEL_REVISION = "6f3ccc0b56e431dc6a0c2b2039706d7d26f22cb9" |
| SFW_REPO = "Phr00t/Qwen-Image-Edit-Rapid-AIO" |
| SFW_FILE = "v19/Qwen-Rapid-AIO-SFW-v19.safetensors" |
| _sfw_path = hf_hub_download(SFW_REPO, SFW_FILE) |
|
|
| _PREFIX = "model.diffusion_model." |
|
|
| _config = QwenImageTransformer2DModel.load_config( |
| BASE_MODEL_ID, subfolder="transformer", revision=BASE_MODEL_REVISION |
| ) |
| with init_empty_weights(): |
| transformer = QwenImageTransformer2DModel.from_config(_config) |
|
|
| _placed = set() |
| with safe_open(_sfw_path, framework="pt", device="cpu") as _f: |
| for _key in _f.keys(): |
| if not _key.startswith(_PREFIX): |
| continue |
| _name = _key[len(_PREFIX):] |
| if _name.startswith("__"): |
| continue |
| set_module_tensor_to_device( |
| transformer, _name, "cpu", value=_f.get_tensor(_key), dtype=dtype |
| ) |
| _placed.add(_name) |
|
|
| _expected = set(transformer.state_dict().keys()) |
| _absent = _expected - _placed |
| if _absent: |
| raise RuntimeError( |
| f"{SFW_FILE} is missing {len(_absent)} of {len(_expected)} tensors, " |
| f"starting with {sorted(_absent)[:5]}" |
| ) |
| transformer.eval() |
| print(f"Loaded {SFW_FILE}: {len(_placed)} tensors", flush=True) |
|
|
| pipe = QwenImageEditPlusPipeline.from_pretrained( |
| BASE_MODEL_ID, |
| revision=BASE_MODEL_REVISION, |
| transformer=transformer, |
| torch_dtype=dtype, |
| ) |
|
|
| del transformer |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| fp8_config = Float8DynamicActivationFloat8WeightConfig() |
| quantize_(pipe.transformer, fp8_config) |
| try: |
| quantize_(pipe.text_encoder, fp8_config) |
| except Exception as exc: |
| print( |
| f"Text-encoder FP8 quantization unavailable ({type(exc).__name__}); " |
| "continuing with the pinned text encoder.", |
| flush=True, |
| ) |
|
|
| pipe = pipe.to(device) |
|
|
| try: |
| pipe.transformer.__class__ = QwenImageTransformer2DModel |
| pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3()) |
| except Exception as exc: |
| print(f"FA3 attention processor unavailable; continuing with default attention: {exc}") |
|
|
|
|
| MAX_SEED = np.iinfo(np.int32).max |
| MIN_INPUT_SIDE = 256 |
|
|
| def use_output_as_input(output_images): |
| if not output_images: |
| return gr.update() |
| first = output_images[0] |
| path = first[0] if isinstance(first, (list, tuple)) else first |
| return gr.update(value=path) |
|
|
|
|
| @spaces.GPU(duration=45) |
| def infer( |
| image_1, |
| image_2, |
| prompt, |
| seed=42, |
| randomize_seed=False, |
| true_guidance_scale=1.0, |
| num_inference_steps=4, |
| height=None, |
| width=None, |
| num_images_per_prompt=1, |
| progress=gr.Progress(track_tqdm=True), |
| ): |
| negative_prompt = " " |
|
|
| if randomize_seed: |
| seed = random.randint(0, MAX_SEED) |
|
|
| generator = torch.Generator(device=device).manual_seed(seed) |
|
|
| pil_images = [] |
| for img in (image_1, image_2): |
| if img is None: |
| continue |
| try: |
| if isinstance(img, str): |
| pil_images.append(Image.open(img).convert("RGB")) |
| elif isinstance(img, Image.Image): |
| pil_images.append(img.convert("RGB")) |
| elif hasattr(img, "name"): |
| pil_images.append(Image.open(img.name).convert("RGB")) |
| except Exception: |
| continue |
|
|
| if not pil_images: |
| raise gr.Error("Upload at least one valid input image before editing.") |
|
|
| for img in pil_images: |
| if min(img.size) < MIN_INPUT_SIDE: |
| raise gr.Error( |
| f"Input images must be at least {MIN_INPUT_SIDE}px on each side. " |
| "Very small source images can produce tiled or grid-like artifacts." |
| ) |
|
|
| print( |
| f"Starting generation: seed={seed}, steps={num_inference_steps}, " |
| f"guidance={true_guidance_scale}, size={width}x{height}, inputs={len(pil_images)}", |
| flush=True, |
| ) |
|
|
| images_pil = pipe( |
| image=pil_images if len(pil_images) > 0 else None, |
| prompt=prompt, |
| height=height, |
| width=width, |
| negative_prompt=negative_prompt, |
| num_inference_steps=num_inference_steps, |
| generator=generator, |
| true_cfg_scale=true_guidance_scale, |
| num_images_per_prompt=num_images_per_prompt, |
| ).images |
|
|
| return images_pil, seed, gr.update(visible=True) |
|
|
|
|
| css = """ |
| #col-container { |
| margin: 0 auto; |
| max-width: 1024px; |
| } |
| #logo-title { |
| text-align: center; |
| } |
| #logo-title img { |
| width: 400px; |
| } |
| #edit_text{margin-top: -62px !important} |
| """ |
|
|
| with gr.Blocks(css=css, delete_cache=(3600, 86400)) as demo: |
| with gr.Column(elem_id="col-container"): |
| gr.HTML(""" |
| <div style="text-align:center;padding:4rem 1rem; |
| font-family:system-ui,sans-serif;color:#999;font-size:.9rem"> |
| This Space has no public interface. |
| </div> |
| """) |
| with gr.Row(visible=False): |
| with gr.Column(): |
| with gr.Row(): |
| image_1 = gr.Image(label="Image 1", type="filepath", interactive=True) |
| image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True) |
|
|
| prompt = gr.Text( |
| label="Prompt 🪄", |
| show_label=True, |
| placeholder="Enter your prompt here...", |
| ) |
| run_button = gr.Button("Edit!", variant="primary") |
| |
| with gr.Accordion("Advanced Settings", open=False): |
| |
| |
| seed = gr.Slider( |
| label="Seed", |
| minimum=0, |
| maximum=MAX_SEED, |
| step=1, |
| value=0, |
| ) |
| |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) |
| |
| with gr.Row(): |
| |
| true_guidance_scale = gr.Slider( |
| label="True guidance scale", |
| minimum=1.0, |
| maximum=10.0, |
| step=0.1, |
| value=1.0 |
| ) |
|
|
| num_inference_steps = gr.Slider( |
| label="Number of inference steps", |
| minimum=1, |
| maximum=40, |
| step=1, |
| value=4, |
| ) |
| |
| with gr.Column(): |
| result = gr.Gallery(label="Result", show_label=False, type="filepath") |
| with gr.Row(): |
| use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False) |
|
|
| with gr.Row(visible=False): |
| gr.Markdown("### 📜 History") |
| clear_history_button = gr.Button("🗑️ Clear History", size="sm", variant="stop") |
| |
| history_gallery = gr.Gallery( |
| label="Click any image to use as input", |
| interactive=False, |
| show_label=True, |
| visible=False |
| ) |
|
|
|
|
| |
|
|
| gr.on( |
| triggers=[run_button.click, prompt.submit], |
| fn=infer, |
| inputs=[ |
| image_1, |
| image_2, |
| prompt, |
| seed, |
| randomize_seed, |
| true_guidance_scale, |
| num_inference_steps, |
| ], |
| outputs=[result, seed, use_output_btn], |
|
|
| ).then( |
| fn=update_history, |
| inputs=[result, history_gallery], |
| outputs=history_gallery, |
|
|
| ) |
|
|
| use_output_btn.click( |
| fn=use_output_as_input, |
| inputs=[result], |
| outputs=[image_1] |
| ) |
|
|
| history_gallery.select( |
| fn=use_history_as_input, |
| inputs=None, |
| outputs=[image_1], |
|
|
| ) |
| |
| clear_history_button.click( |
| fn=lambda: [], |
| inputs=None, |
| outputs=history_gallery, |
|
|
| ) |
|
|
|
|
| def _dressfit_load(src): |
| import base64 as _b64, urllib.request as _url |
| from io import BytesIO as _BytesIO |
| if src.startswith("http://") or src.startswith("https://"): |
| req = _url.Request(src, headers={"User-Agent": "dressfit/1.0"}) |
| with _url.urlopen(req, timeout=30) as resp: |
| raw = resp.read(25 * 1024 * 1024) |
| else: |
| if src.startswith("data:"): |
| src = src.split(",", 1)[1] |
| raw = _b64.b64decode(src) |
| return Image.open(_BytesIO(raw)).convert("RGB") |
|
|
|
|
| def dressfit_tryon(person: str, garment: str, prompt: str, |
| steps: float = 4, guidance: float = 1.0, |
| seed: float = 0, randomize: bool = True) -> dict: |
| import base64 as _b64, traceback as _tb |
| from io import BytesIO as _BytesIO |
| try: |
| result, used_seed = infer( |
| _dressfit_load(person), _dressfit_load(garment), prompt, |
| int(seed), bool(randomize), float(guidance), int(steps), |
| )[:2] |
| if isinstance(result, (list, tuple)): |
| result = result[0] if result else None |
| if result is None: |
| return {"status": "error", "message": "the pipeline returned nothing"} |
| buf = _BytesIO() |
| result.save(buf, format="PNG") |
| return {"status": "success", "seed": int(used_seed), |
| "image": "data:image/png;base64," + _b64.b64encode(buf.getvalue()).decode()} |
| except Exception as exc: |
| return {"status": "error", "message": "".join(_tb.format_exception(exc))[-2500:]} |
|
|
|
|
| with demo: |
| gr.api(dressfit_tryon, api_name="dressfit_tryon") |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
