Initial commit for ZeroGPU Space deployment

README.md

@@ -0,0 +1,15 @@
+---
+title: FireRed Image Edit 1.0 Fast
+emoji: 🔥
+colorFrom: red
+colorTo: yellow
+sdk: gradio
+sdk_version: 6.20.0
+python_version: '3.12'
+app_file: app.py
+pinned: true
+license: apache-2.0
+short_description: FireRed-Image-Edit × Qwen-Image-Edit-Rapid (Transformers)
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,1234 @@
+import os
+import gc
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+import base64
+import json
+import html as html_lib
+from io import BytesIO
+from PIL import Image
+
+MAX_SEED = np.iinfo(np.int32).max
+LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__ =", torch.__version__)
+print("torch.version.cuda =", torch.version.cuda)
+print("cuda available:", torch.cuda.is_available())
+print("cuda device count:", torch.cuda.device_count())
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+
+print("Using device:", device)
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+
+dtype = torch.bfloat16
+
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    "FireRedTeam/FireRed-Image-Edit-1.1",
+    transformer=QwenImageTransformer2DModel.from_pretrained(
+        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
+        torch_dtype=dtype,
+        device_map="cuda",
+    ),
+    torch_dtype=dtype,
+).to(device)
+
+try:
+    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+    print("Flash Attention 3 Processor set successfully.")
+except Exception as e:
+    print(f"Warning: Could not set FA3 processor: {e}")
+
+EXAMPLES_CONFIG = []
+
+
+def make_thumb_b64(path, max_dim=220):
+    if not os.path.exists(path):
+        return ""
+    try:
+        img = Image.open(path).convert("RGB")
+        img.thumbnail((max_dim, max_dim), LANCZOS)
+        buf = BytesIO()
+        img.save(buf, format="JPEG", quality=65)
+        return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
+    except Exception as e:
+        print(f"Thumbnail error for {path}: {e}")
+        return ""
+
+
+def encode_full_image(path):
+    if not os.path.exists(path):
+        return ""
+    try:
+        with open(path, "rb") as f:
+            data = f.read()
+        ext = path.rsplit(".", 1)[-1].lower()
+        mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
+        return f"data:{mime};base64,{base64.b64encode(data).decode()}"
+    except Exception as e:
+        print(f"Encode error for {path}: {e}")
+        return ""
+
+
+def build_example_cards_html():
+    cards = ""
+    for i, ex in enumerate(EXAMPLES_CONFIG):
+        thumbs_html = ""
+        for path in ex["images"]:
+            thumb = make_thumb_b64(path)
+            if thumb:
+                thumbs_html += f'<img src="{thumb}" alt="">'
+            else:
+                thumbs_html += '<div class="example-thumb-placeholder">Preview</div>'
+        n = len(ex["images"])
+        badge = f'{n} image{"s" if n > 1 else ""}'
+        prompt_short = html_lib.escape(ex["prompt"][:90])
+        if len(ex["prompt"]) > 90:
+            prompt_short += "..."
+        cards += f'''<div class="example-card" data-idx="{i}">
+            <div class="example-thumbs">{thumbs_html}</div>
+            <div class="example-meta"><span class="example-badge">{badge}</span></div>
+            <div class="example-prompt-text">{prompt_short}</div>
+        </div>'''
+    return cards
+
+
+def load_example_data(idx_str):
+    try:
+        idx = int(float(idx_str)) if idx_str and idx_str.strip() else -1
+    except (ValueError, TypeError):
+        idx = -1
+    if idx < 0 or idx >= len(EXAMPLES_CONFIG):
+        return json.dumps({"images": [], "prompt": "", "names": [], "status": "error"})
+    ex = EXAMPLES_CONFIG[idx]
+    b64_list, names = [], []
+    for path in ex["images"]:
+        b64 = encode_full_image(path)
+        if b64:
+            b64_list.append(b64)
+            names.append(os.path.basename(path))
+    return json.dumps({"images": b64_list, "prompt": ex["prompt"], "names": names, "status": "ok"})
+
+
+print("Building example thumbnails...")
+EXAMPLE_CARDS_HTML = build_example_cards_html()
+print(f"Built {len(EXAMPLES_CONFIG)} example cards.")
+
+
+def b64_to_pil_list(b64_json_str):
+    if not b64_json_str or b64_json_str.strip() in ("", "[]"):
+        return []
+    try:
+        b64_list = json.loads(b64_json_str)
+    except Exception:
+        return []
+    pil_images = []
+    for b64_str in b64_list:
+        if not b64_str or not isinstance(b64_str, str):
+            continue
+        try:
+            if b64_str.startswith("data:image"):
+                _, data = b64_str.split(",", 1)
+            else:
+                data = b64_str
+            image_data = base64.b64decode(data)
+            pil_images.append(Image.open(BytesIO(image_data)).convert("RGB"))
+        except Exception as e:
+            print(f"Error decoding image: {e}")
+    return pil_images
+
+
+def update_dimensions_on_upload(image):
+    if image is None:
+        return 1024, 1024
+    w, h = image.size
+    if w > h:
+        nw = 1024
+        nh = int(nw * h / w)
+    else:
+        nh = 1024
+        nw = int(nh * w / h)
+    return (nw // 8) * 8, (nh // 8) * 8
+
+
+@spaces.GPU(size="xlarge")
+def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+    gc.collect()
+    torch.cuda.empty_cache()
+    pil_images = b64_to_pil_list(images_b64_json)
+    if not pil_images:
+        raise gr.Error("Please upload at least one image to edit.")
+    if not prompt or prompt.strip() == "":
+        raise gr.Error("Please enter an edit prompt.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
+    width, height = update_dimensions_on_upload(pil_images[0])
+    try:
+        result_image = pipe(
+            image=pil_images, prompt=prompt, negative_prompt=negative_prompt,
+            height=height, width=width, num_inference_steps=steps,
+            generator=generator, true_cfg_scale=guidance_scale,
+        ).images[0]
+        return result_image, seed
+    except Exception as e:
+        raise e
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+css = r"""
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=JetBrains+Mono:wght@400;500;600&display=swap');
+*{box-sizing:border-box;margin:0;padding:0}
+body,.gradio-container{
+    background:#0f0f13!important;font-family:'Inter',system-ui,-apple-system,sans-serif!important;
+    font-size:14px!important;color:#e4e4e7!important;min-height:100vh;
+}
+.dark body,.dark .gradio-container{background:#0f0f13!important;color:#e4e4e7!important}
+footer{display:none!important}
+.hidden-input{display:none!important;height:0!important;overflow:hidden!important;margin:0!important;padding:0!important}
+
+#example-load-btn{
+    position:absolute!important;left:-9999px!important;top:-9999px!important;
+    width:1px!important;height:1px!important;opacity:0.01!important;
+    pointer-events:none!important;overflow:hidden!important;
+}
+#gradio-run-btn{
+    position:absolute;left:-9999px;top:-9999px;width:1px;height:1px;
+    opacity:0.01;pointer-events:none;overflow:hidden;
+}
+
+/* ── App shell ── */
+.app-shell{
+    background:#18181b;border:1px solid #27272a;border-radius:16px;
+    margin:12px auto;max-width:1400px;overflow:hidden;
+    box-shadow:0 25px 50px -12px rgba(0,0,0,.6),0 0 0 1px rgba(255,255,255,.03);
+}
+
+/* ── Header ── */
+.app-header{
+    background:linear-gradient(135deg,#18181b,#1e1e24);border-bottom:1px solid #27272a;
+    padding:14px 24px;display:flex;align-items:center;justify-content:space-between;
+    flex-wrap:wrap;gap:12px;
+}
+.app-header-left{display:flex;align-items:center;gap:12px}
+.app-logo{
+    width:36px;height:36px;background:linear-gradient(135deg,#FF0000,#FF3333,#FF8080);
+    border-radius:10px;display:flex;align-items:center;justify-content:center;
+    box-shadow:0 4px 12px rgba(255,0,0,.35);flex-shrink:0;
+}
+.app-logo svg{width:20px;height:20px;fill:#fff;flex-shrink:0}
+.app-title{
+    font-size:18px;font-weight:700;background:linear-gradient(135deg,#e4e4e7,#a1a1aa);
+    -webkit-background-clip:text;-webkit-text-fill-color:transparent;letter-spacing:-.3px;
+}
+.app-badge{
+    font-size:11px;font-weight:600;padding:3px 10px;border-radius:20px;
+    background:rgba(255,0,0,.15);color:#FF3333;border:1px solid rgba(255,0,0,.25);letter-spacing:.3px;
+}
+.app-badge.fast{background:rgba(34,197,94,.12);color:#4ade80;border:1px solid rgba(34,197,94,.25)}
+
+/* ── GitHub button ── */
+.gh-btn{
+    display:inline-flex!important;align-items:center!important;gap:7px!important;
+    padding:7px 16px!important;border-radius:8px!important;text-decoration:none!important;
+    font-family:'Inter',sans-serif!important;font-size:13px!important;font-weight:700!important;
+    letter-spacing:.1px!important;background:#FF0000!important;
+    color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;
+    border:1px solid rgba(255,255,255,.18)!important;
+    box-shadow:0 2px 10px rgba(255,0,0,.45),0 1px 0 rgba(255,255,255,.1) inset!important;
+    transition:transform .15s ease,box-shadow .15s ease,background .15s ease!important;
+    cursor:pointer!important;flex-shrink:0!important;
+}
+.gh-btn:hover{
+    background:#FF3333!important;color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;
+    transform:translateY(-1px)!important;
+    box-shadow:0 5px 18px rgba(255,0,0,.6),0 1px 0 rgba(255,255,255,.12) inset!important;
+}
+.gh-btn:active{
+    background:#CC0000!important;transform:translateY(0)!important;
+    box-shadow:0 1px 5px rgba(255,0,0,.35)!important;
+}
+.gh-btn svg{fill:#ffffff!important;flex-shrink:0;width:15px!important;height:15px!important}
+.gh-btn span{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important}
+
+/* ── Toolbar ── */
+.app-toolbar{
+    background:#18181b;border-bottom:1px solid #27272a;padding:8px 16px;
+    display:flex;gap:4px;align-items:center;flex-wrap:wrap;
+}
+.tb-sep{width:1px;height:28px;background:#27272a;margin:0 8px}
+.modern-tb-btn{
+    display:inline-flex;align-items:center;justify-content:center;gap:6px;
+    min-width:32px;height:34px;background:transparent;border:1px solid transparent;
+    border-radius:8px;cursor:pointer;font-size:13px;font-weight:600;padding:0 12px;
+    font-family:'Inter',sans-serif;color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;
+    transition:all .15s ease;
+}
+.modern-tb-btn:hover{background:rgba(255,0,0,.15);border-color:rgba(255,0,0,.3)}
+.modern-tb-btn:active,.modern-tb-btn.active{background:rgba(255,0,0,.25);border-color:rgba(255,0,0,.45)}
+.modern-tb-btn .tb-label{font-size:13px;color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;font-weight:600}
+.modern-tb-btn .tb-svg{width:15px;height:15px;flex-shrink:0;color:#ffffff!important}
+.modern-tb-btn .tb-svg,
+.modern-tb-btn .tb-svg *{stroke:#ffffff!important;fill:none!important}
+.tb-info{font-family:'JetBrains Mono',monospace;font-size:12px;color:#71717a;padding:0 8px;display:flex;align-items:center}
+
+body:not(.dark) .modern-tb-btn,body:not(.dark) .modern-tb-btn *{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important}
+body:not(.dark) .modern-tb-btn .tb-svg,body:not(.dark) .modern-tb-btn .tb-svg *{stroke:#ffffff!important}
+.dark .modern-tb-btn,.dark .modern-tb-btn *{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important}
+.dark .modern-tb-btn .tb-svg,.dark .modern-tb-btn .tb-svg *{stroke:#ffffff!important}
+.gradio-container .modern-tb-btn,.gradio-container .modern-tb-btn *{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important}
+.gradio-container .modern-tb-btn .tb-svg,.gradio-container .modern-tb-btn .tb-svg *{stroke:#ffffff!important}
+
+/* ── Main layout ── */
+.app-main-row{display:flex;gap:0;flex:1;overflow:hidden}
+.app-main-left{flex:1;display:flex;flex-direction:column;min-width:0;border-right:1px solid #27272a}
+.app-main-right{width:420px;display:flex;flex-direction:column;flex-shrink:0;background:#18181b}
+
+/* ── Drop zone ── */
+#gallery-drop-zone{position:relative;background:#09090b;min-height:440px;overflow:auto}
+#gallery-drop-zone.drag-over{outline:2px solid #FF0000;outline-offset:-2px;background:rgba(255,0,0,.04)}
+
+.upload-prompt-modern{position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);z-index:20}
+.upload-click-area{
+    display:flex;flex-direction:column;align-items:center;justify-content:center;
+    cursor:pointer;padding:36px 52px;border:2px dashed #3f3f46;border-radius:16px;
+    background:rgba(255,0,0,.03);transition:all .2s ease;gap:8px;
+}
+.upload-click-area:hover{background:rgba(255,0,0,.08);border-color:#FF0000;transform:scale(1.03)}
+.upload-click-area:active{background:rgba(255,0,0,.12);transform:scale(.98)}
+.upload-click-area svg{width:80px;height:80px}
+.upload-main-text{color:#71717a;font-size:14px;font-weight:500;margin-top:4px}
+.upload-sub-text{color:#52525b;font-size:12px;text-align:center;max-width:280px;line-height:1.5}
+
+/* ── Gallery grid ── */
+.image-gallery-grid{
+    display:grid;grid-template-columns:repeat(auto-fill,minmax(140px,1fr));
+    gap:12px;padding:16px;align-content:start;
+}
+.gallery-thumb{
+    position:relative;aspect-ratio:1;border-radius:10px;overflow:hidden;
+    cursor:pointer;border:2px solid #27272a;transition:all .2s ease;background:#18181b;
+}
+.gallery-thumb:hover{border-color:#3f3f46;transform:translateY(-2px);box-shadow:0 4px 12px rgba(0,0,0,.4)}
+.gallery-thumb.selected{border-color:#FF0000!important;box-shadow:0 0 0 3px rgba(255,0,0,.2)}
+.gallery-thumb img{width:100%;height:100%;object-fit:cover}
+.thumb-badge{
+    position:absolute;top:6px;left:6px;background:#FF0000;color:#fff;
+    padding:2px 8px;border-radius:4px;font-family:'JetBrains Mono',monospace;font-size:11px;font-weight:600;
+}
+.thumb-remove{
+    position:absolute;top:6px;right:6px;width:24px;height:24px;background:rgba(0,0,0,.75);
+    color:#fff;border:1px solid rgba(255,255,255,.15);border-radius:50%;cursor:pointer;
+    display:none;align-items:center;justify-content:center;font-size:12px;transition:all .15s;line-height:1;
+}
+.gallery-thumb:hover .thumb-remove{display:flex}
+.thumb-remove:hover{background:#FF0000;border-color:#FF0000}
+.gallery-add-card{
+    aspect-ratio:1;border-radius:10px;border:2px dashed #3f3f46;
+    display:flex;flex-direction:column;align-items:center;justify-content:center;
+    cursor:pointer;transition:all .2s ease;background:rgba(255,0,0,.03);gap:4px;
+}
+.gallery-add-card:hover{border-color:#FF0000;background:rgba(255,0,0,.08)}
+.gallery-add-card .add-icon{font-size:28px;color:#71717a;font-weight:300}
+.gallery-add-card .add-text{font-size:12px;color:#71717a;font-weight:500}
+
+/* ── Hint bar ── */
+.hint-bar{
+    background:rgba(255,0,0,.06);border-top:1px solid #27272a;border-bottom:1px solid #27272a;
+    padding:10px 20px;font-size:13px;color:#a1a1aa;line-height:1.7;
+}
+.hint-bar b{color:#FF8080;font-weight:600}
+.hint-bar kbd{
+    display:inline-block;padding:1px 6px;background:#27272a;border:1px solid #3f3f46;
+    border-radius:4px;font-family:'JetBrains Mono',monospace;font-size:11px;color:#a1a1aa;
+}
+
+/* ── Suggestions ── */
+.suggestions-section{border-top:1px solid #27272a;padding:12px 16px}
+.suggestions-title,.examples-title{
+    font-size:12px;font-weight:600;color:#71717a;text-transform:uppercase;
+    letter-spacing:.8px;margin-bottom:10px;
+}
+.suggestions-wrap{display:flex;flex-wrap:wrap;gap:6px}
+.suggestion-chip{
+    display:inline-flex;align-items:center;gap:4px;padding:5px 12px;
+    background:rgba(255,0,0,.08);border:1px solid rgba(255,0,0,.2);border-radius:20px;
+    color:#FF8080;font-size:12px;font-weight:500;font-family:'Inter',sans-serif;
+    cursor:pointer;transition:all .15s;white-space:nowrap;
+}
+.suggestion-chip:hover{background:rgba(255,0,0,.15);border-color:rgba(255,0,0,.35);color:#FF3333;transform:translateY(-1px)}
+
+/* ── Examples ── */
+.examples-section{border-top:1px solid #27272a;padding:12px 16px}
+.examples-scroll{display:flex;gap:10px;overflow-x:auto;padding-bottom:8px}
+.examples-scroll::-webkit-scrollbar{height:6px}
+.examples-scroll::-webkit-scrollbar-track{background:#09090b;border-radius:3px}
+.examples-scroll::-webkit-scrollbar-thumb{background:#27272a;border-radius:3px}
+.examples-scroll::-webkit-scrollbar-thumb:hover{background:#3f3f46}
+.example-card{
+    flex-shrink:0;width:210px;background:#09090b;border:1px solid #27272a;
+    border-radius:10px;overflow:hidden;cursor:pointer;transition:all .2s ease;
+}
+.example-card:hover{border-color:#FF0000;transform:translateY(-2px);box-shadow:0 4px 12px rgba(255,0,0,.15)}
+.example-card.loading{opacity:.5;pointer-events:none}
+.example-thumbs{display:flex;height:110px;overflow:hidden;background:#18181b}
+.example-thumbs img{flex:1;object-fit:cover;min-width:0;border-bottom:1px solid #27272a}
+.example-thumb-placeholder{
+    flex:1;display:flex;align-items:center;justify-content:center;
+    background:#18181b;color:#3f3f46;font-size:11px;min-width:0;
+}
+.example-meta{padding:6px 10px;display:flex;align-items:center;gap:6px}
+.example-badge{
+    display:inline-flex;padding:2px 7px;background:rgba(255,0,0,.1);border-radius:4px;
+    font-size:10px;font-weight:600;color:#FF3333;font-family:'JetBrains Mono',monospace;white-space:nowrap;
+}
+.example-prompt-text{
+    padding:0 10px 8px;font-size:11px;color:#a1a1aa;line-height:1.4;
+    display:-webkit-box;-webkit-line-clamp:2;-webkit-box-orient:vertical;overflow:hidden;
+}
+
+/* ── Right panel ── */
+.panel-card{border-bottom:1px solid #27272a}
+.panel-card-title{
+    padding:12px 20px;font-size:12px;font-weight:600;color:#71717a;
+    text-transform:uppercase;letter-spacing:.8px;border-bottom:1px solid rgba(39,39,42,.6);
+}
+.panel-card-body{padding:16px 20px;display:flex;flex-direction:column;gap:8px}
+.modern-label{font-size:13px;font-weight:500;color:#a1a1aa;margin-bottom:4px;display:block}
+.modern-textarea{
+    width:100%;background:#09090b;border:1px solid #27272a;border-radius:8px;
+    padding:10px 14px;font-family:'Inter',sans-serif;font-size:14px;color:#e4e4e7;
+    resize:vertical;outline:none;min-height:42px;transition:border-color .2s;
+}
+.modern-textarea:focus{border-color:#FF0000;box-shadow:0 0 0 3px rgba(255,0,0,.15)}
+.modern-textarea::placeholder{color:#3f3f46}
+.modern-textarea.error-flash{
+    border-color:#ef4444!important;box-shadow:0 0 0 3px rgba(239,68,68,.2)!important;animation:shake .4s ease;
+}
+@keyframes shake{0%,100%{transform:translateX(0)}20%,60%{transform:translateX(-4px)}40%,80%{transform:translateX(4px)}}
+
+/* ── Toast ── */
+.toast-notification{
+    position:fixed;top:24px;left:50%;transform:translateX(-50%) translateY(-120%);
+    z-index:9999;padding:10px 24px;border-radius:10px;font-family:'Inter',sans-serif;
+    font-size:14px;font-weight:600;display:flex;align-items:center;gap:8px;
+    box-shadow:0 8px 24px rgba(0,0,0,.5);
+    transition:transform .35s cubic-bezier(.34,1.56,.64,1),opacity .35s ease;opacity:0;pointer-events:none;
+}
+.toast-notification.visible{transform:translateX(-50%) translateY(0);opacity:1;pointer-events:auto}
+.toast-notification.error{background:linear-gradient(135deg,#dc2626,#b91c1c);color:#fff;border:1px solid rgba(255,255,255,.15)}
+.toast-notification.warning{background:linear-gradient(135deg,#d97706,#b45309);color:#fff;border:1px solid rgba(255,255,255,.15)}
+.toast-notification.info{background:linear-gradient(135deg,#2563eb,#1d4ed8);color:#fff;border:1px solid rgba(255,255,255,.15)}
+.toast-notification .toast-icon{font-size:16px;line-height:1}
+.toast-notification .toast-text{line-height:1.3}
+
+/* ── Run button ── */
+.btn-run{
+    display:flex;align-items:center;justify-content:center;gap:8px;width:100%;
+    background:linear-gradient(135deg,#FF0000,#CC0000);border:none;border-radius:10px;
+    padding:12px 24px;cursor:pointer;font-size:15px;font-weight:600;font-family:'Inter',sans-serif;
+    color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;transition:all .2s ease;letter-spacing:-.2px;
+    box-shadow:0 4px 16px rgba(255,0,0,.3),inset 0 1px 0 rgba(255,255,255,.1);
+}
+.btn-run:hover{
+    background:linear-gradient(135deg,#FF3333,#FF0000);transform:translateY(-1px);
+    box-shadow:0 6px 24px rgba(255,0,0,.45),inset 0 1px 0 rgba(255,255,255,.15);
+}
+.btn-run:active{transform:translateY(0);box-shadow:0 2px 8px rgba(255,0,0,.3)}
+.btn-run svg{width:18px;height:18px;fill:#ffffff!important}
+.btn-run svg path{fill:#ffffff!important}
+#custom-run-btn,#custom-run-btn *,#custom-run-btn span,#custom-run-btn svg,
+#custom-run-btn svg path,#run-btn-label,.btn-run,.btn-run *{
+    color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;fill:#ffffff!important;
+}
+body:not(.dark) .btn-run,body:not(.dark) .btn-run *,body:not(.dark) #custom-run-btn,
+body:not(.dark) #custom-run-btn *{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;fill:#ffffff!important}
+.dark .btn-run,.dark .btn-run *,.dark #custom-run-btn,.dark #custom-run-btn *{
+    color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;fill:#ffffff!important;
+}
+.gradio-container .btn-run,.gradio-container .btn-run *,.gradio-container #custom-run-btn,
+.gradio-container #custom-run-btn *{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important;fill:#ffffff!important}
+
+/* ── Output ── */
+.output-frame{border-bottom:1px solid #27272a;display:flex;flex-direction:column;position:relative}
+.output-frame .out-title{
+    padding:10px 20px;font-size:13px;font-weight:700;color:#ffffff!important;
+    -webkit-text-fill-color:#ffffff!important;text-transform:uppercase;letter-spacing:.8px;
+    border-bottom:1px solid rgba(39,39,42,.6);display:flex;align-items:center;justify-content:space-between;
+}
+.output-frame .out-title span{color:#ffffff!important;-webkit-text-fill-color:#ffffff!important}
+.output-frame .out-body{
+    flex:1;background:#09090b;display:flex;align-items:center;justify-content:center;
+    overflow:hidden;min-height:240px;position:relative;
+}
+.output-frame .out-body img{max-width:100%;max-height:460px;image-rendering:auto}
+.output-frame .out-placeholder{color:#3f3f46;font-size:13px;text-align:center;padding:20px}
+.out-download-btn{
+    display:none;align-items:center;justify-content:center;background:rgba(255,0,0,.1);
+    border:1px solid rgba(255,0,0,.2);border-radius:6px;cursor:pointer;padding:3px 10px;
+    font-size:11px;font-weight:500;color:#FF8080!important;gap:4px;height:24px;transition:all .15s;
+}
+.out-download-btn:hover{background:rgba(255,0,0,.2);border-color:rgba(255,0,0,.35);color:#ffffff!important}
+.out-download-btn.visible{display:inline-flex}
+.out-download-btn svg{width:12px;height:12px;fill:#FF8080}
+
+/* ── Loader ── */
+.modern-loader{
+    display:none;position:absolute;top:0;left:0;right:0;bottom:0;background:rgba(9,9,11,.92);
+    z-index:15;flex-direction:column;align-items:center;justify-content:center;gap:16px;backdrop-filter:blur(4px);
+}
+.modern-loader.active{display:flex}
+.modern-loader .loader-spinner{
+    width:36px;height:36px;border:3px solid #27272a;border-top-color:#FF0000;
+    border-radius:50%;animation:spin .8s linear infinite;
+}
+@keyframes spin{to{transform:rotate(360deg)}}
+.modern-loader .loader-text{font-size:13px;color:#a1a1aa;font-weight:500}
+.loader-bar-track{width:200px;height:4px;background:#27272a;border-radius:2px;overflow:hidden}
+.loader-bar-fill{
+    height:100%;background:linear-gradient(90deg,#FF0000,#FF3333,#FF0000);
+    background-size:200% 100%;animation:shimmer 1.5s ease-in-out infinite;border-radius:2px;
+}
+@keyframes shimmer{0%{background-position:200% 0}100%{background-position:-200% 0}}
+
+/* ── Settings ── */
+.settings-group{border:1px solid #27272a;border-radius:10px;margin:12px 16px;padding:0;overflow:hidden}
+.settings-group-title{
+    font-size:12px;font-weight:600;color:#71717a;text-transform:uppercase;letter-spacing:.8px;
+    padding:10px 16px;border-bottom:1px solid #27272a;background:rgba(24,24,27,.5);
+}
+.settings-group-body{padding:14px 16px;display:flex;flex-direction:column;gap:12px}
+.slider-row{display:flex;align-items:center;gap:10px;min-height:28px}
+.slider-row label{font-size:13px;font-weight:500;color:#a1a1aa;min-width:72px;flex-shrink:0}
+.slider-row input[type="range"]{
+    flex:1;-webkit-appearance:none;appearance:none;height:6px;background:#27272a;
+    border-radius:3px;outline:none;min-width:0;
+}
+.slider-row input[type="range"]::-webkit-slider-thumb{
+    -webkit-appearance:none;width:16px;height:16px;background:linear-gradient(135deg,#FF0000,#CC0000);
+    border-radius:50%;cursor:pointer;box-shadow:0 2px 6px rgba(255,0,0,.4);transition:transform .15s;
+}
+.slider-row input[type="range"]::-webkit-slider-thumb:hover{transform:scale(1.2)}
+.slider-row input[type="range"]::-moz-range-thumb{
+    width:16px;height:16px;background:linear-gradient(135deg,#FF0000,#CC0000);
+    border-radius:50%;cursor:pointer;border:none;box-shadow:0 2px 6px rgba(255,0,0,.4);
+}
+.slider-row .slider-val{
+    min-width:52px;text-align:right;font-family:'JetBrains Mono',monospace;font-size:12px;
+    font-weight:500;padding:3px 8px;background:#09090b;border:1px solid #27272a;
+    border-radius:6px;color:#a1a1aa;flex-shrink:0;
+}
+.checkbox-row{display:flex;align-items:center;gap:8px;font-size:13px;color:#a1a1aa}
+.checkbox-row input[type="checkbox"]{accent-color:#FF0000;width:16px;height:16px;cursor:pointer}
+.checkbox-row label{color:#a1a1aa;font-size:13px;cursor:pointer}
+
+/* ── Status bar ── */
+.app-statusbar{
+    background:#18181b;border-top:1px solid #27272a;padding:6px 20px;
+    display:flex;gap:12px;height:34px;align-items:center;font-size:12px;
+}
+.app-statusbar .sb-section{
+    padding:0 12px;flex:1;display:flex;align-items:center;font-family:'JetBrains Mono',monospace;
+    font-size:12px;color:#52525b;overflow:hidden;white-space:nowrap;
+}
+.app-statusbar .sb-section.sb-fixed{
+    flex:0 0 auto;min-width:90px;text-align:center;justify-content:center;
+    padding:3px 12px;background:rgba(255,0,0,.08);border-radius:6px;color:#FF3333;font-weight:500;
+}
+
+/* ── Footer note ── */
+.exp-note{
+    padding:10px 20px;font-size:12px;color:#52525b;
+    border-top:1px solid #27272a;text-align:center;font-weight:500;
+    background:#18181b;font-family:'Inter',sans-serif;
+}
+.exp-note a{color:#FF3333;text-decoration:none}
+.exp-note a:hover{text-decoration:underline}
+
+/* ── Dark overrides ── */
+.dark .app-shell{background:#18181b}
+.dark .upload-prompt-modern{background:transparent}
+.dark .panel-card{background:#18181b}
+.dark .settings-group{background:#18181b}
+.dark .output-frame .out-title{color:#ffffff!important}
+.dark .output-frame .out-title span{color:#ffffff!important}
+.dark .out-download-btn{color:#FF8080!important}
+.dark .out-download-btn:hover{color:#ffffff!important}
+
+/* ── Scrollbars ── */
+::-webkit-scrollbar{width:8px;height:8px}
+::-webkit-scrollbar-track{background:#09090b}
+::-webkit-scrollbar-thumb{background:#27272a;border-radius:4px}
+::-webkit-scrollbar-thumb:hover{background:#3f3f46}
+
+/* ── Responsive ── */
+@media(max-width:840px){
+    .app-main-row{flex-direction:column}
+    .app-main-right{width:100%}
+    .app-main-left{border-right:none;border-bottom:1px solid #27272a}
+}
+"""
+
+gallery_js = r"""
+() => {
+function init() {
+    if (window.__fireRedInitDone) return;
+
+    const galleryGrid   = document.getElementById('image-gallery-grid');
+    const dropZone      = document.getElementById('gallery-drop-zone');
+    const uploadPrompt  = document.getElementById('upload-prompt');
+    const uploadClick   = document.getElementById('upload-click-area');
+    const fileInput     = document.getElementById('custom-file-input');
+    const btnUpload     = document.getElementById('tb-upload');
+    const btnRemove     = document.getElementById('tb-remove');
+    const btnClear      = document.getElementById('tb-clear');
+    const promptInput   = document.getElementById('custom-prompt-input');
+    const runBtnEl      = document.getElementById('custom-run-btn');
+    const imgCountTb    = document.getElementById('tb-image-count');
+    const imgCountSb    = document.getElementById('sb-image-count');
+
+    if (!galleryGrid || !fileInput || !dropZone) {
+        setTimeout(init, 250);
+        return;
+    }
+
+    window.__fireRedInitDone = true;
+
+    let images = [];
+    window.__uploadedImages = images;
+    let selectedIdx = -1;
+    let toastTimer = null;
+
+    /* ── GitHub button hover ── */
+    function enforceGhBtn() {
+        const ghBtn = document.querySelector('.gh-btn');
+        if (ghBtn && !ghBtn.__hoverBound) {
+            ghBtn.__hoverBound = true;
+            ghBtn.addEventListener('mouseenter', () => {
+                ghBtn.style.setProperty('background','#FF3333','important');
+                ghBtn.style.setProperty('transform','translateY(-1px)','important');
+                ghBtn.style.setProperty('box-shadow','0 5px 18px rgba(255,0,0,.6)','important');
+            });
+            ghBtn.addEventListener('mouseleave', () => {
+                ghBtn.style.setProperty('background','#FF0000','important');
+                ghBtn.style.setProperty('transform','translateY(0)','important');
+                ghBtn.style.setProperty('box-shadow','0 2px 10px rgba(255,0,0,.45)','important');
+            });
+            ghBtn.addEventListener('mousedown', () => ghBtn.style.setProperty('background','#CC0000','important'));
+            ghBtn.addEventListener('mouseup',   () => ghBtn.style.setProperty('background','#FF3333','important'));
+        }
+    }
+    enforceGhBtn();
+    setInterval(enforceGhBtn, 1000);
+
+    function showToast(message, type) {
+        let toast = document.getElementById('app-toast');
+        if (!toast) {
+            toast = document.createElement('div');
+            toast.id = 'app-toast';
+            toast.className = 'toast-notification';
+            toast.innerHTML = '<span class="toast-icon"></span><span class="toast-text"></span>';
+            document.body.appendChild(toast);
+        }
+        const icon = toast.querySelector('.toast-icon');
+        const text = toast.querySelector('.toast-text');
+        toast.className = 'toast-notification ' + (type || 'error');
+        if (type === 'warning') icon.textContent = '\u26A0';
+        else if (type === 'info') icon.textContent = '\u2139';
+        else icon.textContent = '\u2717';
+        text.textContent = message;
+        if (toastTimer) clearTimeout(toastTimer);
+        void toast.offsetWidth;
+        toast.classList.add('visible');
+        toastTimer = setTimeout(() => toast.classList.remove('visible'), 3500);
+    }
+    window.__showToast = showToast;
+
+    function flashPromptError() {
+        if (!promptInput) return;
+        promptInput.classList.add('error-flash');
+        promptInput.focus();
+        setTimeout(() => promptInput.classList.remove('error-flash'), 800);
+    }
+
+    function setGradioValue(containerId, value) {
+        const container = document.getElementById(containerId);
+        if (!container) return;
+        container.querySelectorAll('input, textarea').forEach(el => {
+            if (el.type === 'file' || el.type === 'range' || el.type === 'checkbox') return;
+            const proto = el.tagName === 'TEXTAREA' ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
+            const ns = Object.getOwnPropertyDescriptor(proto, 'value');
+            if (ns && ns.set) {
+                ns.set.call(el, value);
+                el.dispatchEvent(new Event('input',  {bubbles:true, composed:true}));
+                el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
+            }
+        });
+    }
+    window.__setGradioValue = setGradioValue;
+
+    function syncImagesToGradio() {
+        window.__uploadedImages = images;
+        const b64Array = images.map(img => img.b64);
+        setGradioValue('hidden-images-b64', JSON.stringify(b64Array));
+        updateCounts();
+    }
+
+    function syncPromptToGradio() {
+        if (promptInput) setGradioValue('prompt-gradio-input', promptInput.value);
+    }
+
+    function updateCounts() {
+        const n = images.length;
+        const txt = n > 0 ? n + ' image' + (n > 1 ? 's' : '') : 'No images';
+        if (imgCountTb) imgCountTb.textContent = txt;
+        if (imgCountSb) imgCountSb.textContent = n > 0 ? txt + ' uploaded' : 'No images uploaded';
+    }
+
+    function addImage(b64, name) {
+        images.push({id: Date.now() + Math.random(), b64: b64, name: name});
+        renderGallery();
+        syncImagesToGradio();
+    }
+    window.__addImage = addImage;
+
+    function removeImage(idx) {
+        images.splice(idx, 1);
+        if (selectedIdx === idx) selectedIdx = -1;
+        else if (selectedIdx > idx) selectedIdx--;
+        renderGallery();
+        syncImagesToGradio();
+    }
+
+    function clearAll() {
+        images = [];
+        window.__uploadedImages = images;
+        selectedIdx = -1;
+        renderGallery();
+        syncImagesToGradio();
+    }
+    window.__clearAll = clearAll;
+
+    function selectImage(idx) {
+        selectedIdx = (selectedIdx === idx) ? -1 : idx;
+        renderGallery();
+    }
+
+    function renderGallery() {
+        if (images.length === 0) {
+            galleryGrid.innerHTML = '';
+            galleryGrid.style.display = 'none';
+            if (uploadPrompt) uploadPrompt.style.display = '';
+            return;
+        }
+        if (uploadPrompt) uploadPrompt.style.display = 'none';
+        galleryGrid.style.display = 'grid';
+
+        let html = '';
+        images.forEach((img, i) => {
+            const sel = i === selectedIdx ? ' selected' : '';
+            html += '<div class="gallery-thumb' + sel + '" data-idx="' + i + '">'
+                  + '<img src="' + img.b64 + '" alt="' + (img.name||'image') + '">'
+                  + '<span class="thumb-badge">#' + (i+1) + '</span>'
+                  + '<button class="thumb-remove" data-remove="' + i + '">\u2715</button>'
+                  + '</div>';
+        });
+        html += '<div class="gallery-add-card" id="gallery-add-card">'
+              + '<span class="add-icon">+</span>'
+              + '<span class="add-text">Add</span>'
+              + '</div>';
+        galleryGrid.innerHTML = html;
+
+        galleryGrid.querySelectorAll('.gallery-thumb').forEach(thumb => {
+            thumb.addEventListener('click', (e) => {
+                if (e.target.closest('.thumb-remove')) return;
+                selectImage(parseInt(thumb.dataset.idx));
+            });
+        });
+        galleryGrid.querySelectorAll('.thumb-remove').forEach(btn => {
+            btn.addEventListener('click', (e) => {
+                e.stopPropagation();
+                removeImage(parseInt(btn.dataset.remove));
+            });
+        });
+        const addCard = document.getElementById('gallery-add-card');
+        if (addCard) addCard.addEventListener('click', () => fileInput.click());
+    }
+
+    function processFiles(files) {
+        Array.from(files).forEach(file => {
+            if (!file.type.startsWith('image/')) return;
+            const reader = new FileReader();
+            reader.onload = (e) => addImage(e.target.result, file.name);
+            reader.readAsDataURL(file);
+        });
+    }
+
+    fileInput.addEventListener('change', (e) => { processFiles(e.target.files); e.target.value = ''; });
+    if (uploadClick) uploadClick.addEventListener('click', () => fileInput.click());
+    if (btnUpload) btnUpload.addEventListener('click', () => fileInput.click());
+    if (btnRemove) btnRemove.addEventListener('click', () => {
+        if (selectedIdx >= 0 && selectedIdx < images.length) removeImage(selectedIdx);
+    });
+    if (btnClear) btnClear.addEventListener('click', clearAll);
+
+    dropZone.addEventListener('dragover', (e) => { e.preventDefault(); dropZone.classList.add('drag-over'); });
+    dropZone.addEventListener('dragleave', (e) => { e.preventDefault(); dropZone.classList.remove('drag-over'); });
+    dropZone.addEventListener('drop', (e) => {
+        e.preventDefault(); dropZone.classList.remove('drag-over');
+        if (e.dataTransfer.files.length) processFiles(e.dataTransfer.files);
+    });
+
+    if (promptInput) promptInput.addEventListener('input', syncPromptToGradio);
+
+    window.__setPrompt = function(text) {
+        if (promptInput) { promptInput.value = text; syncPromptToGradio(); }
+    };
+
+    document.querySelectorAll('.example-card[data-idx]').forEach(card => {
+        card.addEventListener('click', () => {
+            const idx = card.getAttribute('data-idx');
+            document.querySelectorAll('.example-card.loading').forEach(c => c.classList.remove('loading'));
+            card.classList.add('loading');
+            showToast('Loading example...', 'info');
+            setGradioValue('example-result-data', '');
+            setGradioValue('example-idx-input', idx);
+            setTimeout(() => {
+                const btn = document.getElementById('example-load-btn');
+                if (btn) {
+                    const b = btn.querySelector('button');
+                    if (b) b.click(); else btn.click();
+                }
+            }, 150);
+            setTimeout(() => card.classList.remove('loading'), 12000);
+        });
+    });
+
+    function syncSlider(customId, gradioId) {
+        const slider = document.getElementById(customId);
+        const valSpan = document.getElementById(customId + '-val');
+        if (!slider) return;
+        slider.addEventListener('input', () => {
+            if (valSpan) valSpan.textContent = slider.value;
+            const container = document.getElementById(gradioId);
+            if (!container) return;
+            container.querySelectorAll('input[type="range"],input[type="number"]').forEach(el => {
+                const ns = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, 'value');
+                if (ns && ns.set) {
+                    ns.set.call(el, slider.value);
+                    el.dispatchEvent(new Event('input',  {bubbles:true, composed:true}));
+                    el.dispatchEvent(new Event('change', {bubbles:true, composed:true}));
+                }
+            });
+        });
+    }
+    syncSlider('custom-seed', 'gradio-seed');
+    syncSlider('custom-guidance', 'gradio-guidance');
+    syncSlider('custom-steps', 'gradio-steps');
+
+    const randCheck = document.getElementById('custom-randomize');
+    if (randCheck) {
+        randCheck.addEventListener('change', () => {
+            const container = document.getElementById('gradio-randomize');
+            if (!container) return;
+            const cb = container.querySelector('input[type="checkbox"]');
+            if (cb && cb.checked !== randCheck.checked) cb.click();
+        });
+    }
+
+    function showLoader() {
+        const l = document.getElementById('output-loader');
+        if (l) l.classList.add('active');
+        const sb = document.querySelector('.sb-fixed');
+        if (sb) sb.textContent = 'Processing...';
+    }
+    function hideLoader() {
+        const l = document.getElementById('output-loader');
+        if (l) l.classList.remove('active');
+        const sb = document.querySelector('.sb-fixed');
+        if (sb) sb.textContent = 'Done';
+    }
+    window.__showLoader = showLoader;
+    window.__hideLoader = hideLoader;
+
+    function validateBeforeRun() {
+        const promptVal = promptInput ? promptInput.value.trim() : '';
+        const hasImages = images.length > 0;
+        if (!hasImages && !promptVal) { showToast('Please upload an image and enter a prompt', 'error'); flashPromptError(); return false; }
+        if (!hasImages) { showToast('Please upload at least one image', 'error'); return false; }
+        if (!promptVal) { showToast('Please enter an edit prompt', 'warning'); flashPromptError(); return false; }
+        return true;
+    }
+
+    window.__clickGradioRunBtn = function() {
+        if (!validateBeforeRun()) return;
+        syncPromptToGradio(); syncImagesToGradio(); showLoader();
+        setTimeout(() => {
+            const gradioBtn = document.getElementById('gradio-run-btn');
+            if (!gradioBtn) return;
+            const btn = gradioBtn.querySelector('button');
+            if (btn) btn.click(); else gradioBtn.click();
+        }, 200);
+    };
+
+    if (runBtnEl) runBtnEl.addEventListener('click', () => window.__clickGradioRunBtn());
+
+    renderGallery();
+    updateCounts();
+}
+init();
+}
+"""
+
+wire_outputs_js = r"""
+() => {
+function watchOutputs() {
+    const resultContainer = document.getElementById('gradio-result');
+    const outBody  = document.getElementById('output-image-container');
+    const outPh    = document.getElementById('output-placeholder');
+    const dlBtn    = document.getElementById('dl-btn-output');
+
+    if (!resultContainer || !outBody) { setTimeout(watchOutputs, 500); return; }
+
+    if (dlBtn) {
+        dlBtn.addEventListener('click', (e) => {
+            e.stopPropagation();
+            const img = outBody.querySelector('img.modern-out-img');
+            if (img && img.src) {
+                const a = document.createElement('a');
+                a.href = img.src; a.download = 'firered_output.png';
+                document.body.appendChild(a); a.click(); document.body.removeChild(a);
+            }
+        });
+    }
+
+    function syncImage() {
+        const resultImg = resultContainer.querySelector('img');
+        if (resultImg && resultImg.src) {
+            if (outPh) outPh.style.display = 'none';
+            let existing = outBody.querySelector('img.modern-out-img');
+            if (!existing) {
+                existing = document.createElement('img');
+                existing.className = 'modern-out-img';
+                outBody.appendChild(existing);
+            }
+            if (existing.src !== resultImg.src) {
+                existing.src = resultImg.src;
+                if (dlBtn) dlBtn.classList.add('visible');
+                if (window.__hideLoader) window.__hideLoader();
+            }
+        }
+    }
+    const observer = new MutationObserver(syncImage);
+    observer.observe(resultContainer, {childList:true, subtree:true, attributes:true, attributeFilter:['src']});
+    setInterval(syncImage, 800);
+}
+watchOutputs();
+
+function watchSeed() {
+    const seedContainer = document.getElementById('gradio-seed');
+    const seedSlider = document.getElementById('custom-seed');
+    const seedVal = document.getElementById('custom-seed-val');
+    if (!seedContainer || !seedSlider) { setTimeout(watchSeed, 500); return; }
+    function sync() {
+        const el = seedContainer.querySelector('input[type="range"],input[type="number"]');
+        if (el && el.value) { seedSlider.value = el.value; if (seedVal) seedVal.textContent = el.value; }
+    }
+    const obs = new MutationObserver(sync);
+    obs.observe(seedContainer, {childList:true, subtree:true, attributes:true, attributeFilter:['value']});
+    setInterval(sync, 1000);
+}
+watchSeed();
+
+function watchExampleResults() {
+    const container = document.getElementById('example-result-data');
+    if (!container) { setTimeout(watchExampleResults, 500); return; }
+
+    let lastProcessed = '';
+
+    function checkResult() {
+        const el = container.querySelector('textarea') || container.querySelector('input');
+        if (!el) return;
+        const val = el.value;
+        if (!val || val === lastProcessed || val.length < 20) return;
+
+        try {
+            const data = JSON.parse(val);
+            if (data.status === 'ok' && data.images && data.images.length > 0) {
+                lastProcessed = val;
+                if (window.__clearAll) window.__clearAll();
+                if (window.__setPrompt && data.prompt) window.__setPrompt(data.prompt);
+                data.images.forEach((b64, i) => {
+                    if (b64 && window.__addImage) {
+                        const name = (data.names && data.names[i]) ? data.names[i] : ('example_' + (i+1) + '.jpg');
+                        window.__addImage(b64, name);
+                    }
+                });
+                document.querySelectorAll('.example-card.loading').forEach(c => c.classList.remove('loading'));
+                if (window.__showToast) window.__showToast('Example loaded \u2014 ' + data.images.length + ' image(s)', 'info');
+            } else if (data.status === 'error') {
+                document.querySelectorAll('.example-card.loading').forEach(c => c.classList.remove('loading'));
+                if (window.__showToast) window.__showToast('Could not load example images', 'error');
+            }
+        } catch(e) {
+            console.error('Example parse error:', e);
+        }
+    }
+
+    const obs = new MutationObserver(checkResult);
+    obs.observe(container, {childList:true, subtree:true, characterData:true, attributes:true});
+    setInterval(checkResult, 500);
+}
+watchExampleResults();
+}
+"""
+
+# ── SVG assets ─────────────────────────────────────────────────────────────────
+DOWNLOAD_SVG = '<svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M12 16l-5-5h3V4h4v7h3l-5 5z"/><path d="M20 18H4v2h16v-2z"/></svg>'
+
+UPLOAD_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4"/><polyline points="17 8 12 3 7 8"/><line x1="12" y1="3" x2="12" y2="15"/></svg>'
+
+REMOVE_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="15" y1="9" x2="9" y2="15"/><line x1="9" y1="9" x2="15" y2="15"/></svg>'
+
+CLEAR_SVG = '<svg class="tb-svg" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="3 6 5 6 21 6"/><path d="M19 6v14a2 2 0 01-2 2H7a2 2 0 01-2-2V6m3 0V4a2 2 0 012-2h4a2 2 0 012 2v2"/><line x1="10" y1="11" x2="10" y2="17"/><line x1="14" y1="11" x2="14" y2="17"/></svg>'
+
+GITHUB_SVG = '<svg width="15" height="15" viewBox="0 0 16 16" xmlns="http://www.w3.org/2000/svg"><path fill="#ffffff" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"/></svg>'
+
+FIRE_LOGO_SVG = '<svg viewBox="0 0 24 24" fill="white" xmlns="http://www.w3.org/2000/svg"><path d="M12 23c-3.6 0-8-2.69-8-7.5 0-3.5 3-6.5 4.5-8 .27-.27.75-.08.75.28v2.44c0 .42.5.63.72.28C12.28 7.5 13 3 13 1c0-.42.48-.64.8-.35C18 4.5 20 9 20 12c0 5.5-3.5 11-8 11z"/></svg>'
+
+# ── Gradio app ─────────────────────────────────────────────────────────────────
+with gr.Blocks() as demo:
+
+    hidden_images_b64 = gr.Textbox(value="[]", elem_id="hidden-images-b64", elem_classes="hidden-input", container=False)
+    prompt = gr.Textbox(value="", elem_id="prompt-gradio-input", elem_classes="hidden-input", container=False)
+    seed = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, elem_id="gradio-seed", elem_classes="hidden-input", container=False)
+    randomize_seed = gr.Checkbox(value=True, elem_id="gradio-randomize", elem_classes="hidden-input", container=False)
+    guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.0, elem_id="gradio-guidance", elem_classes="hidden-input", container=False)
+    steps = gr.Slider(minimum=1, maximum=50, step=1, value=4, elem_id="gradio-steps", elem_classes="hidden-input", container=False)
+    result = gr.Image(elem_id="gradio-result", elem_classes="hidden-input", container=False, format="png")
+
+    example_idx = gr.Textbox(value="", elem_id="example-idx-input", elem_classes="hidden-input", container=False)
+    example_result = gr.Textbox(value="", elem_id="example-result-data", elem_classes="hidden-input", container=False)
+    example_load_btn = gr.Button("Load Example", elem_id="example-load-btn")
+
+    gr.HTML(f"""
+    <div class="app-shell">
+
+        <!-- Header with GitHub top-right -->
+        <div class="app-header">
+            <div class="app-header-left">
+                <div class="app-logo">{FIRE_LOGO_SVG}</div>
+                <span class="app-title">FireRed-Image-Edit</span>
+                <span class="app-badge">v1.1</span>
+                <span class="app-badge fast">4-Step Fast</span>
+            </div>
+            <a href="https://github.com/PRITHIVSAKTHIUR/FireRed-Image-Edit-1.0-Fast"
+               target="_blank" class="gh-btn">
+                {GITHUB_SVG}
+                <span>GitHub</span>
+            </a>
+        </div>
+
+        <!-- Toolbar -->
+        <div class="app-toolbar">
+            <button id="tb-upload" class="modern-tb-btn" title="Upload images">
+                {UPLOAD_SVG}<span class="tb-label">Upload</span>
+            </button>
+            <button id="tb-remove" class="modern-tb-btn" title="Remove selected image">
+                {REMOVE_SVG}<span class="tb-label">Remove</span>
+            </button>
+            <button id="tb-clear" class="modern-tb-btn" title="Clear all images">
+                {CLEAR_SVG}<span class="tb-label">Clear All</span>
+            </button>
+            <div class="tb-sep"></div>
+            <span id="tb-image-count" class="tb-info">No images</span>
+        </div>
+
+        <!-- Main row -->
+        <div class="app-main-row">
+
+            <!-- Left panel -->
+            <div class="app-main-left">
+                <div id="gallery-drop-zone">
+                    <div id="upload-prompt" class="upload-prompt-modern">
+                        <div id="upload-click-area" class="upload-click-area">
+                            <svg viewBox="0 0 80 80" fill="none" xmlns="http://www.w3.org/2000/svg">
+                                <rect x="8" y="14" width="64" height="52" rx="6" fill="none"
+                                      stroke="#FF0000" stroke-width="2" stroke-dasharray="4 3"/>
+                                <polygon points="12,62 30,40 42,50 54,34 68,62"
+                                         fill="rgba(255,0,0,0.15)" stroke="#FF0000" stroke-width="1.5"/>
+                                <circle cx="28" cy="30" r="6"
+                                        fill="rgba(255,0,0,0.2)" stroke="#FF0000" stroke-width="1.5"/>
+                            </svg>
+                            <span class="upload-main-text">Click or drag images here</span>
+                            <span class="upload-sub-text">Supports multiple images for reference-based editing and guided manipulation</span>
+                        </div>
+                    </div>
+                    <input id="custom-file-input" type="file" accept="image/*" multiple style="display:none;" />
+                    <div id="image-gallery-grid" class="image-gallery-grid" style="display:none;"></div>
+                </div>
+
+                <div class="hint-bar">
+                    <b>Upload:</b> Click or drag to add images &nbsp;&middot;&nbsp;
+                    <b>Multi-image:</b> Upload multiple images for reference-based editing &nbsp;&middot;&nbsp;
+                    <kbd>Remove</kbd> deletes selected &nbsp;&middot;&nbsp;
+                    <kbd>Clear All</kbd> removes everything
+                </div>
+
+                <div class="suggestions-section">
+                    <div class="suggestions-title">Quick Prompts</div>
+                    <div class="suggestions-wrap">
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Transform the image into a dotted cartoon style.')">Cartoon Style</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Convert it to black and white.')">Black and White</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Add cinematic lighting with warm orange tones and film grain.')">Cinematic</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Transform into anime style illustration.')">Anime Style</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Apply oil painting effect with visible brush strokes.')">Oil Painting</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Enhance and upscale with more detail and clarity.')">Enhance</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Make it look like a watercolor painting with soft edges.')">Watercolor</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Add dramatic sunset sky and warm lighting.')">Sunset Glow</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Convert to detailed pencil sketch with cross-hatching and shading.')">Pencil Sketch</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Apply pop art style with bold colors and halftone patterns.')">Pop Art</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Apply a vintage retro film look with faded colors and light leaks.')">Vintage Retro</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Add neon glow effects with vibrant colors against a dark background.')">Neon Glow</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Convert to pixel art style with a retro 16-bit aesthetic.')">Pixel Art</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Simplify into a clean minimalist illustration with flat colors.')">Minimalist</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Convert to low poly 3D geometric art style.')">Low Poly 3D</button>
+                        <button class="suggestion-chip" onclick="window.__setPrompt('Transform into comic book style with bold outlines and cel shading.')">Comic Book</button>
+                    </div>
+                </div>
+
+                <div class="examples-section">
+                    <div class="examples-title">Quick Examples &mdash; click to load</div>
+                    <div class="examples-scroll">
+                        {EXAMPLE_CARDS_HTML}
+                    </div>
+                </div>
+            </div>
+
+            <!-- Right panel -->
+            <div class="app-main-right">
+                <div class="panel-card">
+                    <div class="panel-card-title">Edit Instruction</div>
+                    <div class="panel-card-body">
+                        <label class="modern-label" for="custom-prompt-input">Prompt</label>
+                        <textarea id="custom-prompt-input" class="modern-textarea" rows="3"
+                                  placeholder="e.g., transform into anime, upscale, change lighting..."></textarea>
+                    </div>
+                </div>
+
+                <div style="padding:12px 20px;">
+                    <button id="custom-run-btn" class="btn-run">
+                        <svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" width="18" height="18">
+                            <path d="M12 23c-3.6 0-8-2.69-8-7.5 0-3.5 3-6.5 4.5-8 .27-.27.75-.08.75.28v2.44c0 .42.5.63.72.28C12.28 7.5 13 3 13 1c0-.42.48-.64.8-.35C18 4.5 20 9 20 12c0 5.5-3.5 11-8 11z" fill="white"/>
+                        </svg>
+                        <span id="run-btn-label">Edit Image</span>
+                    </button>
+                </div>
+
+                <div class="output-frame" style="flex:1">
+                    <div class="out-title">
+                        <span>Output</span>
+                        <span id="dl-btn-output" class="out-download-btn" title="Download">
+                            {DOWNLOAD_SVG} Save
+                        </span>
+                    </div>
+                    <div class="out-body" id="output-image-container">
+                        <div class="modern-loader" id="output-loader">
+                            <div class="loader-spinner"></div>
+                            <div class="loader-text">Processing image...</div>
+                            <div class="loader-bar-track"><div class="loader-bar-fill"></div></div>
+                        </div>
+                        <div class="out-placeholder" id="output-placeholder">Result will appear here</div>
+                    </div>
+                </div>
+
+                <div class="settings-group">
+                    <div class="settings-group-title">Advanced Settings</div>
+                    <div class="settings-group-body">
+                        <div class="slider-row">
+                            <label>Seed</label>
+                            <input type="range" id="custom-seed" min="0" max="2147483647" step="1" value="0">
+                            <span class="slider-val" id="custom-seed-val">0</span>
+                        </div>
+                        <div class="checkbox-row">
+                            <input type="checkbox" id="custom-randomize" checked>
+                            <label for="custom-randomize">Randomize seed</label>
+                        </div>
+                        <div class="slider-row">
+                            <label>Guidance</label>
+                            <input type="range" id="custom-guidance" min="1" max="10" step="0.1" value="1.0">
+                            <span class="slider-val" id="custom-guidance-val">1.0</span>
+                        </div>
+                        <div class="slider-row">
+                            <label>Steps</label>
+                            <input type="range" id="custom-steps" min="1" max="50" step="1" value="4">
+                            <span class="slider-val" id="custom-steps-val">4</span>
+                        </div>
+                    </div>
+                </div>
+            </div>
+        </div>
+
+        <!-- Footer: only model credit, no GitHub link -->
+        <div class="exp-note">
+            Experimental Space for
+            <a href="https://huggingface.co/FireRedTeam/FireRed-Image-Edit-1.1" target="_blank">FireRed-Image-Edit-1.1</a>
+        </div>
+
+        <!-- Status bar -->
+        <div class="app-statusbar">
+            <div class="sb-section" id="sb-image-count">No images uploaded</div>
+            <div class="sb-section sb-fixed">Ready</div>
+        </div>
+
+    </div><!-- /app-shell -->
+    """)
+
+    run_btn = gr.Button("Run", elem_id="gradio-run-btn")
+
+    demo.load(fn=None, js=gallery_js)
+    demo.load(fn=None, js=wire_outputs_js)
+
+    run_btn.click(
+        fn=infer,
+        inputs=[hidden_images_b64, prompt, seed, randomize_seed, guidance_scale, steps],
+        outputs=[result, seed],
+        js=r"""(imgs, p, s, rs, gs, st) => {
+            const images = window.__uploadedImages || [];
+            const b64Array = images.map(img => img.b64);
+            const imgsJson = JSON.stringify(b64Array);
+            const promptEl = document.getElementById('custom-prompt-input');
+            const promptVal = promptEl ? promptEl.value : p;
+            return [imgsJson, promptVal, s, rs, gs, st];
+        }""",
+    )
+
+    example_load_btn.click(
+        fn=load_example_data,
+        inputs=[example_idx],
+        outputs=[example_result],
+        queue=False,
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=50).launch(
+        css=css,
+        mcp_server=True,
+        ssr_mode=False,
+        show_error=True,
+        allowed_paths=["examples"],
+    )

pre-requirements.txt

@@ -0,0 +1 @@
+pip==26.1.2

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -0,0 +1,891 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import math
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import QwenImageLoraLoaderMixin
+from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from PIL import Image
+        >>> from diffusers import QwenImageEditPlusPipeline
+        >>> from diffusers.utils import load_image
+
+        >>> pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> image = load_image(
+        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
+        ... ).convert("RGB")
+        >>> prompt = (
+        ...     "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors"
+        ... )
+        >>> # Depending on the variant being used, the pipeline call will slightly vary.
+        >>> # Refer to the pipeline documentation for more details.
+        >>> image = pipe(image, prompt, num_inference_steps=50).images[0]
+        >>> image.save("qwenimage_edit_plus.png")
+        ```
+"""
+
+CONDITION_IMAGE_SIZE = 384 * 384
+VAE_IMAGE_SIZE = 1024 * 1024
+
+
+# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def calculate_dimensions(target_area, ratio):
+    width = math.sqrt(target_area * ratio)
+    height = width / ratio
+
+    width = round(width / 32) * 32
+    height = round(height / 32) * 32
+
+    return width, height
+
+
+class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
+    r"""
+    The Qwen-Image-Edit pipeline for image editing.
+
+    Args:
+        transformer ([`QwenImageTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant.
+        tokenizer (`QwenTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLQwenImage,
+        text_encoder: Qwen2_5_VLForConditionalGeneration,
+        tokenizer: Qwen2Tokenizer,
+        processor: Qwen2VLProcessor,
+        transformer: QwenImageTransformer2DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            processor=processor,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
+        # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = 1024
+
+        self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 64
+        self.default_sample_size = 128
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
+    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
+        bool_mask = mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+        return split_result
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        image: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+        if isinstance(image, list):
+            base_img_prompt = ""
+            for i, img in enumerate(image):
+                base_img_prompt += img_prompt_template.format(i + 1)
+        elif image is not None:
+            base_img_prompt = img_prompt_template.format(1)
+        else:
+            base_img_prompt = ""
+
+        template = self.prompt_template_encode
+
+        drop_idx = self.prompt_template_encode_start_idx
+        txt = [template.format(base_img_prompt + e) for e in prompt]
+
+        model_inputs = self.processor(
+            text=txt,
+            images=image,
+            padding=True,
+            return_tensors="pt",
+        ).to(device)
+
+        outputs = self.text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=model_inputs.pixel_values,
+            image_grid_thw=model_inputs.image_grid_thw,
+            output_hidden_states=True,
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        return prompt_embeds, encoder_attention_mask
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        image: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 1024,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            image (`torch.Tensor`, *optional*):
+                image to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds_mask=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_embeds_mask is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 1024:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+
+    @staticmethod
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._unpack_latents
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
+
+        return latents
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        image_latents = (image_latents - latents_mean) / latents_std
+
+        return image_latents
+
+    def prepare_latents(
+        self,
+        images,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, 1, num_channels_latents, height, width)
+
+        image_latents = None
+        if images is not None:
+            if not isinstance(images, list):
+                images = [images]
+            all_image_latents = []
+            for image in images:
+                image = image.to(device=device, dtype=dtype)
+                if image.shape[1] != self.latent_channels:
+                    image_latents = self._encode_vae_image(image=image, generator=generator)
+                else:
+                    image_latents = image
+                if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+                    # expand init_latents for batch_size
+                    additional_image_per_prompt = batch_size // image_latents.shape[0]
+                    image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+                elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+                    raise ValueError(
+                        f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+                    )
+                else:
+                    image_latents = torch.cat([image_latents], dim=0)
+
+                image_latent_height, image_latent_width = image_latents.shape[3:]
+                image_latents = self._pack_latents(
+                    image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
+                )
+                all_image_latents.append(image_latents)
+            image_latents = torch.cat(all_image_latents, dim=1)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        return latents, image_latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: Optional[PipelineImageInput] = None,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        true_cfg_scale: float = 4.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: Optional[float] = None,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
+                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
+                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
+                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
+                latents as `image`, but if passing latents directly it is not encoded again.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                true_cfg_scale (`float`, *optional*, defaults to 1.0): Guidance scale as defined in [Classifier-Free
+                Diffusion Guidance](https://huggingface.co/papers/2207.12598). `true_cfg_scale` is defined as `w` of
+                equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Classifier-free guidance is
+                enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale
+                encourages to generate images that are closely linked to the text `prompt`, usually at the expense of
+                lower image quality.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to None):
+                A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance
+                where the guidance scale is applied during inference through noise prediction rescaling, guidance
+                distilled models take the guidance scale directly as an input parameter during forward pass. Guidance
+                scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
+                that are closely linked to the text `prompt`, usually at the expense of lower image quality. This
+                parameter in the pipeline is there to support future guidance-distilled models when they come up. It is
+                ignored when not using guidance distilled models. To enable traditional classifier-free guidance,
+                please pass `true_cfg_scale > 1.0` and `negative_prompt` (even an empty negative prompt like " " should
+                enable classifier-free guidance computations).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`:
+            [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is a list with the generated images.
+        """
+        image_size = image[-1].size if isinstance(image, list) else image.size
+        calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
+        height = height or calculated_height
+        width = width or calculated_width
+
+        multiple_of = self.vae_scale_factor * 2
+        width = width // multiple_of * multiple_of
+        height = height // multiple_of * multiple_of
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # 3. Preprocess image
+        if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
+            if not isinstance(image, list):
+                image = [image]
+            condition_image_sizes = []
+            condition_images = []
+            vae_image_sizes = []
+            vae_images = []
+            for img in image:
+                image_width, image_height = img.size
+                condition_width, condition_height = calculate_dimensions(
+                    CONDITION_IMAGE_SIZE, image_width / image_height
+                )
+                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
+                condition_image_sizes.append((condition_width, condition_height))
+                vae_image_sizes.append((vae_width, vae_height))
+                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
+                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+
+        if true_cfg_scale > 1 and not has_neg_prompt:
+            logger.warning(
+                f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no negative_prompt is provided."
+            )
+        elif true_cfg_scale <= 1 and has_neg_prompt:
+            logger.warning(
+                " negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1"
+            )
+
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            image=condition_images,
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                image=condition_images,
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, image_latents = self.prepare_latents(
+            vae_images,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        img_shapes = [
+            [
+                (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
+                *[
+                    (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2)
+                    for vae_width, vae_height in vae_image_sizes
+                ],
+            ]
+        ] * batch_size
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds and guidance_scale is None:
+            raise ValueError("guidance_scale is required for guidance-distilled model.")
+        elif self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        elif not self.transformer.config.guidance_embeds and guidance_scale is not None:
+            logger.warning(
+                f"guidance_scale is passed as {guidance_scale}, but ignored since the model is not guidance-distilled."
+            )
+            guidance = None
+        elif not self.transformer.config.guidance_embeds and guidance_scale is None:
+            guidance = None
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        
+        image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=latents.device)
+        if do_true_cfg:
+            negative_txt_seq_lens = (
+                negative_prompt_embeds_mask.sum(dim=1).tolist()
+                if negative_prompt_embeds_mask is not None
+                else None
+            )
+            uncond_image_rotary_emb = self.transformer.pos_embed(
+                img_shapes, negative_txt_seq_lens, device=latents.device
+            )
+        else:
+            uncond_image_rotary_emb = None
+
+        # 6. Denoising loop
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                with self.transformer.cache_context("cond"):
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        encoder_hidden_states_mask=prompt_embeds_mask,
+                        encoder_hidden_states=prompt_embeds,
+                        image_rotary_emb=image_rotary_emb,
+                        attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred[:, : latents.size(1)]
+
+                if do_true_cfg:
+                    with self.transformer.cache_context("uncond"):
+                        neg_noise_pred = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep / 1000,
+                            guidance=guidance,
+                            encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            image_rotary_emb=uncond_image_rotary_emb,
+                            attention_kwargs=self.attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                    noise_pred = comb_pred * (cond_norm / noise_norm)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return QwenImagePipelineOutput(images=image)

qwenimage/qwen_fa3_processor.py

@@ -0,0 +1,233 @@
+"""
+Paired with a good language model. Thanks!
+
+FA3 is currently broken on Blackwell (sm_100) GPUs; this module detects that
+at import time and falls back to PyTorch scaled-dot-product attention (SDPA)
+automatically.  The public class name / call signature are unchanged.
+"""
+
+import torch
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
+
+
+# ---------------------------------------------------------------------------
+# FA3 availability check
+# ---------------------------------------------------------------------------
+
+def _is_blackwell() -> bool:
+    """Return True when the current default CUDA device is an sm_100 (Blackwell) GPU."""
+    if not torch.cuda.is_available():
+        return False
+    cap = torch.cuda.get_device_capability()
+    # Blackwell → compute capability 10.x  (sm_100)
+    return cap[0] >= 10
+
+
+_fa3_available: bool = False
+_fa3_unavailable_reason: str = ""
+_flash_attn_func = None
+
+if _is_blackwell():
+    _fa3_unavailable_reason = (
+        "FlashAttention-3 is not yet supported on Blackwell (sm_100) GPUs. "
+        "Falling back to scaled-dot-product attention (SDPA)."
+    )
+else:
+    try:
+        from kernels import get_kernel
+        _k = get_kernel("kernels-community/vllm-flash-attn3")
+        _flash_attn_func = _k.flash_attn_func
+        _fa3_available = True
+    except Exception as e:
+        _fa3_unavailable_reason = (
+            "FlashAttention-3 via Hugging Face `kernels` is unavailable. "
+            f"Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n{e}\n"
+            "Falling back to scaled-dot-product attention (SDPA)."
+        )
+
+
+# ---------------------------------------------------------------------------
+# FA3 custom op (registered only when the kernel loaded successfully)
+# ---------------------------------------------------------------------------
+
+if _fa3_available:
+    @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
+    def flash_attn_func(
+        q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
+    ) -> torch.Tensor:
+        # _flash_attn_func returns (output, softmax_lse); we only need output.
+        output, _lse = _flash_attn_func(q, k, v, causal=causal)
+        return output
+
+    @flash_attn_func.register_fake
+    def _flash_attn_func_fake(q, k, v, causal=False):
+        # output shape mirrors q: (batch, seq_len, num_heads, head_dim)
+        return torch.empty_like(q).contiguous()
+
+else:
+    # Provide a stub so call-sites that import the symbol don't break at
+    # module load; the processor will route around it at runtime.
+    def flash_attn_func(
+        q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
+    ) -> torch.Tensor:
+        raise RuntimeError(_fa3_unavailable_reason)
+
+
+# ---------------------------------------------------------------------------
+# SDPA fallback helper
+# ---------------------------------------------------------------------------
+
+def _sdpa_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    causal: bool = False,
+) -> torch.Tensor:
+    """
+    Scaled dot-product attention using torch.nn.functional.scaled_dot_product_attention.
+
+    Input / output layout: (B, S, H, D_h)  — same as the FA3 kernel.
+    """
+    # SDPA expects (B, H, S, D_h)
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+
+    out = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
+
+    # Back to (B, S, H, D_h)
+    return out.transpose(1, 2)
+
+
+# ---------------------------------------------------------------------------
+# Attention processor
+# ---------------------------------------------------------------------------
+
+class QwenDoubleStreamAttnProcessorFA3:
+    """
+    Attention processor for the Qwen double-stream architecture.
+
+    Preferred backend: vLLM FlashAttention-3 via Hugging Face ``kernels``.
+    Automatic fallback: PyTorch ``scaled_dot_product_attention`` (SDPA) when
+    FA3 is unavailable — e.g. on Blackwell (sm_100) GPUs where FA3 is not yet
+    supported, or when the ``kernels`` package is absent.
+
+    Notes / limitations
+    -------------------
+    - Arbitrary attention masks are not supported on the FA3 path.  Pass
+      ``attention_mask=None`` (the default) to stay on the fast path.
+    - On the SDPA path, ``attention_mask`` is likewise ignored; add explicit
+      support here if you need it.
+    - ``encoder_hidden_states`` (text stream) is required.
+    """
+
+    _attention_backend: str  # set in __init__ after capability detection
+
+    def __init__(self):
+        if _fa3_available:
+            self._attention_backend = "fa3"
+        else:
+            import warnings
+            warnings.warn(
+                f"QwenDoubleStreamAttnProcessorFA3: {_fa3_unavailable_reason}",
+                stacklevel=2,
+            )
+            self._attention_backend = "sdpa"
+
+    def _attend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        causal: bool = False,
+    ) -> torch.Tensor:
+        """Dispatch to FA3 or SDPA depending on what is available."""
+        if self._attention_backend == "fa3":
+            return flash_attn_func(q, k, v, causal=causal)
+        return _sdpa_attention(q, k, v, causal=causal)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.FloatTensor,                          # (B, S_img, D_model)
+        encoder_hidden_states: torch.FloatTensor = None,           # (B, S_txt, D_model)
+        encoder_hidden_states_mask: torch.FloatTensor = None,      # unused
+        attention_mask: Optional[torch.FloatTensor] = None,        # unsupported on FA3 path
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+
+        if encoder_hidden_states is None:
+            raise ValueError(
+                "QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream)."
+            )
+        if attention_mask is not None and self._attention_backend == "fa3":
+            raise NotImplementedError(
+                "attention_mask is not supported on the FA3 path. "
+                "Either drop the mask or let the processor fall back to SDPA."
+            )
+
+        B, S_img, _ = hidden_states.shape
+        S_txt = encoder_hidden_states.shape[1]
+
+        # ---- QKV projections ----
+        img_q = attn.to_q(hidden_states)
+        img_k = attn.to_k(hidden_states)
+        img_v = attn.to_v(hidden_states)
+
+        txt_q = attn.add_q_proj(encoder_hidden_states)
+        txt_k = attn.add_k_proj(encoder_hidden_states)
+        txt_v = attn.add_v_proj(encoder_hidden_states)
+
+        # ---- Reshape to (B, S, H, D_h) ----
+        H = attn.heads
+        img_q = img_q.unflatten(-1, (H, -1))
+        img_k = img_k.unflatten(-1, (H, -1))
+        img_v = img_v.unflatten(-1, (H, -1))
+
+        txt_q = txt_q.unflatten(-1, (H, -1))
+        txt_k = txt_k.unflatten(-1, (H, -1))
+        txt_v = txt_v.unflatten(-1, (H, -1))
+
+        # ---- Q/K normalization ----
+        if getattr(attn, "norm_q", None) is not None:
+            img_q = attn.norm_q(img_q)
+        if getattr(attn, "norm_k", None) is not None:
+            img_k = attn.norm_k(img_k)
+        if getattr(attn, "norm_added_q", None) is not None:
+            txt_q = attn.norm_added_q(txt_q)
+        if getattr(attn, "norm_added_k", None) is not None:
+            txt_k = attn.norm_added_k(txt_k)
+
+        # ---- RoPE (Qwen variant) ----
+        if image_rotary_emb is not None:
+            img_freqs, txt_freqs = image_rotary_emb
+            img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
+            img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
+            txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
+            txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
+
+        # ---- Joint attention over [text, image] along sequence axis ----
+        q = torch.cat([txt_q, img_q], dim=1)  # (B, S_txt + S_img, H, D_h)
+        k = torch.cat([txt_k, img_k], dim=1)
+        v = torch.cat([txt_v, img_v], dim=1)
+
+        out = self._attend(q, k, v, causal=False)  # (B, S_total, H, D_h)
+
+        # ---- Back to (B, S, D_model) ----
+        out = out.flatten(2, 3).to(q.dtype)
+
+        # ---- Split text / image segments ----
+        txt_attn_out = out[:, :S_txt, :]
+        img_attn_out = out[:, S_txt:, :]
+
+        # ---- Output projections ----
+        img_attn_out = attn.to_out[0](img_attn_out)
+        if len(attn.to_out) > 1:
+            img_attn_out = attn.to_out[1](img_attn_out)  # dropout if present
+
+        txt_attn_out = attn.to_add_out(txt_attn_out)
+
+        return img_attn_out, txt_attn_out

qwenimage/transformer_qwenimage.py

@@ -0,0 +1,642 @@
+# Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import FeedForward, AttentionMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.attention_processor import Attention
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    Args
+        timesteps (torch.Tensor):
+            a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        embedding_dim (int):
+            the dimension of the output.
+        flip_sin_to_cos (bool):
+            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+        downscale_freq_shift (float):
+            Controls the delta between frequencies between dimensions
+        scale (float):
+            Scaling factor applied to the embeddings.
+        max_period (int):
+            Controls the maximum frequency of the embeddings
+    Returns
+        torch.Tensor: an [N x dim] Tensor of positional embeddings.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+def apply_rotary_emb_qwen(
+    x: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(1)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)
+
+
+class QwenTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timestep, hidden_states):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
+
+        conditioning = timesteps_emb
+
+        return conditioning
+
+
+class QwenEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [
+                self.rope_params(pos_index, self.axes_dim[0], self.theta),
+                self.rope_params(pos_index, self.axes_dim[1], self.theta),
+                self.rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.neg_freqs = torch.cat(
+            [
+                self.rope_params(neg_index, self.axes_dim[0], self.theta),
+                self.rope_params(neg_index, self.axes_dim[1], self.theta),
+                self.rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.rope_cache = {}
+
+        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+        self.scale_rope = scale_rope
+
+    def rope_params(self, index, dim, theta=10000):
+        """
+        Args:
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+        """
+        assert dim % 2 == 0
+        freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)))
+        freqs = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs
+
+    def forward(self, video_fhw, txt_seq_lens, device):
+        """
+        Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
+        txt_length: [bs] a list of 1 integers representing the length of the text
+        """
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_freqs = []
+        max_vid_index = 0
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            rope_key = f"{idx}_{height}_{width}"
+
+            if not torch.compiler.is_compiling():
+                if rope_key not in self.rope_cache:
+                    self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
+                video_freq = self.rope_cache[rope_key]
+            else:
+                video_freq = self._compute_video_freqs(frame, height, width, idx)
+            video_freq = video_freq.to(device)
+            vid_freqs.append(video_freq)
+
+            if self.scale_rope:
+                max_vid_index = max(height // 2, width // 2, max_vid_index)
+            else:
+                max_vid_index = max(height, width, max_vid_index)
+
+        max_len = max(txt_seq_lens)
+        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        return vid_freqs, txt_freqs
+
+    @functools.lru_cache(maxsize=None)
+    def _compute_video_freqs(self, frame, height, width, idx=0):
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+        if self.scale_rope:
+            freqs_height = torch.cat([freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0)
+            freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0)
+            freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+        else:
+            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+        return freqs.clone().contiguous()
+
+
+class QwenDoubleStreamAttnProcessor2_0:
+    """
+    Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
+    implements joint attention computation where text and image streams are processed together.
+    """
+
+    _attention_backend = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
+        encoder_hidden_states_mask: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
+
+        seq_txt = encoder_hidden_states.shape[1]
+
+        # Compute QKV for image stream (sample projections)
+        img_query = attn.to_q(hidden_states)
+        img_key = attn.to_k(hidden_states)
+        img_value = attn.to_v(hidden_states)
+
+        # Compute QKV for text stream (context projections)
+        txt_query = attn.add_q_proj(encoder_hidden_states)
+        txt_key = attn.add_k_proj(encoder_hidden_states)
+        txt_value = attn.add_v_proj(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        img_query = img_query.unflatten(-1, (attn.heads, -1))
+        img_key = img_key.unflatten(-1, (attn.heads, -1))
+        img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+        # Apply QK normalization
+        if attn.norm_q is not None:
+            img_query = attn.norm_q(img_query)
+        if attn.norm_k is not None:
+            img_key = attn.norm_k(img_key)
+        if attn.norm_added_q is not None:
+            txt_query = attn.norm_added_q(txt_query)
+        if attn.norm_added_k is not None:
+            txt_key = attn.norm_added_k(txt_key)
+
+        # Apply RoPE
+        if image_rotary_emb is not None:
+            img_freqs, txt_freqs = image_rotary_emb
+            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
+            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
+            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
+            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
+
+        # Concatenate for joint attention
+        # Order: [text, image]
+        joint_query = torch.cat([txt_query, img_query], dim=1)
+        joint_key = torch.cat([txt_key, img_key], dim=1)
+        joint_value = torch.cat([txt_value, img_value], dim=1)
+
+        # Compute joint attention
+        joint_hidden_states = dispatch_attention_fn(
+            joint_query,
+            joint_key,
+            joint_value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=self._attention_backend,
+        )
+
+        # Reshape back
+        joint_hidden_states = joint_hidden_states.flatten(2, 3)
+        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
+
+        # Split attention outputs back
+        txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
+        img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part
+
+        # Apply output projections
+        img_attn_output = attn.to_out[0](img_attn_output)
+        if len(attn.to_out) > 1:
+            img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+        txt_attn_output = attn.to_add_out(txt_attn_output)
+
+        return img_attn_output, txt_attn_output
+
+
+@maybe_allow_in_graph
+class QwenImageTransformerBlock(nn.Module):
+    def __init__(
+        self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+
+        # Image processing modules
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,  # Enable cross attention for joint computation
+            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=QwenDoubleStreamAttnProcessor2_0(),
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # Text processing modules
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        # Text doesn't need separate attention - it's handled by img_attn joint computation
+        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def _modulate(self, x, mod_params):
+        """Apply modulation to input tensor"""
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_mask: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Get modulation parameters for both streams
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # Split modulation parameters for norm1 and norm2
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+        # Process image stream - norm1 + modulation
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+
+        # Process text stream - norm1 + modulation
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+
+        # Use QwenAttnProcessor2_0 for joint attention computation
+        # This directly implements the DoubleStreamLayerMegatron logic:
+        # 1. Computes QKV for both streams
+        # 2. Applies QK normalization and RoPE
+        # 3. Concatenates and runs joint attention
+        # 4. Splits results back to separate streams
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual (like in Megatron)
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        # Process image stream - norm2 + MLP
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        # Process text stream - norm2 + MLP
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+        # Clip to prevent overflow for fp16
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin):
+    """
+    The Transformer model introduced in Qwen.
+
+    Args:
+        patch_size (`int`, defaults to `2`):
+            Patch size to turn the input data into small patches.
+        in_channels (`int`, defaults to `64`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `None`):
+            The number of channels in the output. If not specified, it defaults to `in_channels`.
+        num_layers (`int`, defaults to `60`):
+            The number of layers of dual stream DiT blocks to use.
+        attention_head_dim (`int`, defaults to `128`):
+            The number of dimensions to use for each attention head.
+        num_attention_heads (`int`, defaults to `24`):
+            The number of attention heads to use.
+        joint_attention_dim (`int`, defaults to `3584`):
+            The number of dimensions to use for the joint attention (embedding/channel dimension of
+            `encoder_hidden_states`).
+        guidance_embeds (`bool`, defaults to `False`):
+            Whether to use guidance embeddings for guidance-distilled variant of the model.
+        axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+            The dimensions to use for the rotary positional embeddings.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["QwenImageTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    _repeated_blocks = ["QwenImageTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 64,
+        out_channels: Optional[int] = 16,
+        num_layers: int = 60,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 3584,
+        guidance_embeds: bool = False,  # TODO: this should probably be removed
+        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+    ):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
+
+        self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+
+        self.img_in = nn.Linear(in_channels, self.inner_dim)
+        self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                QwenImageTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        encoder_hidden_states_mask: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        guidance: torch.Tensor = None,  # TODO: this should probably be removed
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        """
+        The [`QwenTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
+                Mask of the input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        hidden_states = self.img_in(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype)
+        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = (
+            self.time_text_embed(timestep, hidden_states)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, hidden_states)
+        )
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_hidden_states_mask,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_states_mask=encoder_hidden_states_mask,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        # Use only the image part (hidden_states) from the dual-stream blocks
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

requirements.txt

@@ -0,0 +1,69 @@
+torch==2.11.0
+torchvision==0.26.0
+transformers==5.14.1
+accelerate==1.14.0
+diffusers==0.39.0
+peft==0.19.1
+tokenizers==0.22.2
+sentencepiece==0.2.2
+safetensors==0.8.0
+gradio==6.20.0
+gradio-client==2.5.0
+hf-gradio==0.4.1
+spaces==0.51.0
+fastapi==0.139.1
+starlette==1.3.1
+uvicorn==0.51.0
+pydantic==2.12.5
+pydantic-core==2.41.5
+pydantic-settings==2.14.2
+typing-inspection==0.4.2
+python-multipart==0.0.32
+orjson==3.11.9
+httpx-sse==0.4.3
+websockets==16.1
+mcp==1.28.1
+platformdirs==4.10.0
+psutil==7.2.2
+regex==2026.7.10
+pillow==12.3.0
+av==18.0.0
+pydub==0.25.1
+authlib==1.7.2
+cryptography==49.0.0
+pyOpenSSL==26.3.0
+cffi==2.1.0
+pycparser==3.0
+email-validator==2.3.0
+dnspython==2.8.0
+python-dotenv==1.2.2
+itsdangerous==2.2.0
+pyjwt==2.13.0
+jsonschema==4.26.0
+jsonschema-specifications==2025.9.1
+referencing==0.37.0
+rpds-py==2026.6.3
+annotated-types==0.7.0
+semantic-version==2.10.0
+tomlkit==0.14.0
+importlib_metadata==9.0.0
+zipp==4.1.0
+pytz==2026.2
+safehttpx==0.1.7
+brotli==1.2.0
+groovy==0.1.2
+id==1.6.1
+joserfc==1.7.3
+kernels==0.16.0
+kernels-data==0.16.0
+rfc3161-client==1.0.7
+rfc8785==0.1.4
+securesystemslib==1.4.0
+sigstore==4.4.0
+sigstore-models==0.0.6
+sigstore-rekor-types==0.0.18
+sse-starlette==3.4.5
+tuf==7.0.0
+
+# scientific computing
+numpy==2.4.6