Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/10.jpg filter=lfs diff=lfs merge=lfs -text
+examples/11.png filter=lfs diff=lfs merge=lfs -text
+examples/6.jpg filter=lfs diff=lfs merge=lfs -text
+examples/7.webp filter=lfs diff=lfs merge=lfs -text
+examples/8.jpg filter=lfs diff=lfs merge=lfs -text
+examples/9.png filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README.md

@@ -1,13 +1,83 @@
 ---
 title: FireRed Image Edit 1.0 Fast
-emoji: 💻
-colorFrom: blue
-colorTo: gray
+emoji: 🔥
+colorFrom: red
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.14.0
-python_version: '3.13'
+sdk_version: "5.0.0"
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# **[FireRed-Image-Edit-1.0-Fast](https://huggingface.co/spaces/prithivMLmods/FireRed-Image-Edit-1.0-Fast)**
+
+FireRed-Image-Edit-1.0-Fast is a high-performance, AI-driven image editing application that utilizes advanced diffusers and the QwenImageEditPlusPipeline for precise, prompt-based image modifications. Incorporating rapid Transformer configurations, the application provides an interactive Gradio web interface with a custom Soft Blue theme for an aesthetically pleasing user experience. Users can leverage powerful flow match euler discrete schedulers to seamlessly edit visual content by submitting an original image alongside descriptive textual instructions. The application operates entirely in Python, efficiently utilizing CUDA capabilities for accelerated machine learning computations, and serves as a fast, state-of-the-art solution for automated, text-guided image manipulation without complex manual editing software.
+
+<img width="1918" height="1753" alt="Screenshot 2026-03-21 at 15-27-14 FireRed Image Edit 1 0 Fast - a Hugging Face Space by prithivMLmods" src="https://github.com/user-attachments/assets/c88a82b6-f877-4312-94e3-fd3119b03318" />
+
+## Features
+
+* **Advanced Diffusers Pipeline:** Utilizes the QwenImageEditPlusPipeline integrated with FlowMatchEulerDiscreteScheduler for high-fidelity image editing based on user prompts.
+* **Rapid AI Architecture:** Employs optimized transformer structures designed for fast inference, providing quick iterations and real-time responsiveness.
+* **Custom Themed Interface:** Provides an interactive, user-friendly Gradio web interface styled with a custom Soft OrangeRed theme for an optimal visual layout.
+* **Hardware Acceleration:** Automatically identifies and leverages CUDA-compatible devices for optimal computational performance, rendering complex edits rapidly.
+
+## Installation
+
+### 1. Clone the Repository
+
+```bash
+git clone https://github.com/PRITHIVSAKTHIUR/FireRed-Image-Edit-1.0-Fast.git
+cd FireRed-Image-Edit-1.0-Fast
+```
+
+### 2. Install Pre-requirements
+
+Certain system-level or structural dependencies must be configured before setting up the main python environment:
+
+```bash
+pip install -r pre-requirements.txt
+```
+
+### 3. Install Standard Dependencies
+
+Install the core Python packages, which include critical modules like Diffusers, Accelerate, PEFT, and Gradio:
+
+```bash
+pip install -r requirements.txt
+```
+
+## How to Run
+
+To start the application and load the local server, run the main Python script:
+
+```bash
+python app.py
+```
+
+Once the model weights are successfully loaded into your device's memory and the server starts, the terminal will provide a local URL (typically `http://127.0.0.1:7860`). Open this link in your web browser to interact with the visual interface.
+
+## Project Structure
+
+* `app.py`: The main entry point script containing the custom Gradio interface setup, pipeline initialization, and inference logic.
+* `qwenimage/`: Core directory housing the transformer and processor modules crucial for the underlying image manipulation techniques.
+* `requirements.txt`: The primary file listing Python library requirements needed to operate the application correctly.
+* `pre-requirements.txt`: A list containing earlier or auxiliary dependency specifications.
+* `examples/`: Directory dedicated to storing sample images and expected outputs to verify application functionality.
+* `LICENSE.txt`: The legal text detailing the licensing constraints and permissions.
+
+## Workflow
+
+1. Navigate to the local server URL provided after executing the application.
+2. Upload a source image that you wish to edit into the input module.
+3. Provide a clear, detailed text prompt describing the exact modifications you want the AI to perform on the image.
+4. The system executes the QwenImageEditPlusPipeline via the underlying rapid transformers to compute the altered visual output.
+5. Retrieve and save the edited image directly from the interface.
+
+## License
+
+This project is open-source. For detailed terms and conditions, refer to the included `LICENSE.txt` file within the repository.
+
+## Contributing
+
+Community contributions are encouraged. Please submit an issue for bug reports or create a Pull Request to propose features, optimize inference times, or improve the user interface.

app.py

@@ -0,0 +1,1158 @@
+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("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 = [
+    {
+        "images": ["examples/1.jpg"],
+        "prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed 'Fire-Edit' preserving realistic texture and details.",
+    },
+    {
+        "images": ["examples/2.jpg"],
+        "prompt": "Transform the image into a dotted cartoon style.",
+    },
+    {
+        "images": ["examples/3.jpeg"],
+        "prompt": "Convert it to black and white.",
+    },
+    {
+        "images": ["examples/4.jpg", "examples/5.jpg"],
+        "prompt": "Replace her glasses with the new glasses from image 1.",
+    },
+    {
+        "images": ["examples/8.jpg", "examples/9.png"],
+        "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
+    },
+    {
+        "images": ["examples/10.jpg", "examples/11.png"],
+        "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
+    },
+]
+
+
+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
+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{
+    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);
+}
+.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,#1E90FF,#47A3FF,#7CB8FF);
+    border-radius:10px;display:flex;align-items:center;justify-content:center;
+    box-shadow:0 4px 12px rgba(30,144,255,.35);
+}
+.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(30,144,255,.15);color:#47A3FF;border:1px solid rgba(30,144,255,.25);letter-spacing:.3px;
+}
+.app-badge.fast{background:rgba(34,197,94,.12);color:#4ade80;border:1px solid rgba(34,197,94,.25)}
+
+.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(30,144,255,.15);border-color:rgba(30,144,255,.3)}
+.modern-tb-btn:active,.modern-tb-btn.active{background:rgba(30,144,255,.25);border-color:rgba(30,144,255,.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}
+
+.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}
+
+#gallery-drop-zone{position:relative;background:#09090b;min-height:440px;overflow:auto}
+#gallery-drop-zone.drag-over{outline:2px solid #1E90FF;outline-offset:-2px;background:rgba(30,144,255,.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(30,144,255,.03);transition:all .2s ease;gap:8px;
+}
+.upload-click-area:hover{background:rgba(30,144,255,.08);border-color:#1E90FF;transform:scale(1.03)}
+.upload-click-area:active{background:rgba(30,144,255,.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}
+
+.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:#1E90FF!important;box-shadow:0 0 0 3px rgba(30,144,255,.2)}
+.gallery-thumb img{width:100%;height:100%;object-fit:cover}
+.thumb-badge{
+    position:absolute;top:6px;left:6px;background:#1E90FF;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:#1E90FF;border-color:#1E90FF}
+.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(30,144,255,.03);gap:4px;
+}
+.gallery-add-card:hover{border-color:#1E90FF;background:rgba(30,144,255,.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{
+    background:rgba(30,144,255,.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:#7CB8FF;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-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(30,144,255,.08);border:1px solid rgba(30,144,255,.2);border-radius:20px;
+    color:#7CB8FF;font-size:12px;font-weight:500;font-family:'Inter',sans-serif;
+    cursor:pointer;transition:all .15s;white-space:nowrap;
+}
+.suggestion-chip:hover{background:rgba(30,144,255,.15);border-color:rgba(30,144,255,.35);color:#47A3FF;transform:translateY(-1px)}
+
+.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:#1E90FF;transform:translateY(-2px);box-shadow:0 4px 12px rgba(30,144,255,.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(30,144,255,.1);border-radius:4px;
+    font-size:10px;font-weight:600;color:#47A3FF;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;
+}
+
+.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:#1E90FF;box-shadow:0 0 0 3px rgba(30,144,255,.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-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}
+
+.btn-run{
+    display:flex;align-items:center;justify-content:center;gap:8px;width:100%;
+    background:linear-gradient(135deg,#1E90FF,#1873CC);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(30,144,255,.3),inset 0 1px 0 rgba(255,255,255,.1);
+}
+.btn-run:hover{
+    background:linear-gradient(135deg,#47A3FF,#1E90FF);transform:translateY(-1px);
+    box-shadow:0 6px 24px rgba(30,144,255,.45),inset 0 1px 0 rgba(255,255,255,.15);
+}
+.btn-run:active{transform:translateY(0);box-shadow:0 2px 8px rgba(30,144,255,.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-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(30,144,255,.1);
+    border:1px solid rgba(30,144,255,.2);border-radius:6px;cursor:pointer;padding:3px 10px;
+    font-size:11px;font-weight:500;color:#7CB8FF!important;gap:4px;height:24px;transition:all .15s;
+}
+.out-download-btn:hover{background:rgba(30,144,255,.2);border-color:rgba(30,144,255,.35);color:#ffffff!important}
+.out-download-btn.visible{display:inline-flex}
+.out-download-btn svg{width:12px;height:12px;fill:#7CB8FF}
+
+.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:#1E90FF;
+    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,#1E90FF,#47A3FF,#1E90FF);
+    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-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,#1E90FF,#1873CC);
+    border-radius:50%;cursor:pointer;box-shadow:0 2px 6px rgba(30,144,255,.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,#1E90FF,#1873CC);
+    border-radius:50%;cursor:pointer;border:none;box-shadow:0 2px 6px rgba(30,144,255,.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:#1E90FF;width:16px;height:16px;cursor:pointer}
+.checkbox-row label{color:#a1a1aa;font-size:13px;cursor:pointer}
+
+.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(30,144,255,.08);border-radius:6px;color:#47A3FF;font-weight:500;
+}
+
+.exp-note{padding:10px 20px;font-size:12px;color:#52525b;border-top:1px solid #27272a;text-align:center}
+.exp-note a{color:#47A3FF;text-decoration:none}
+.exp-note a:hover{text-decoration:underline}
+
+.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:#7CB8FF!important}
+.dark .out-download-btn:hover{color:#ffffff!important}
+
+::-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}
+
+@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;
+
+    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 — ' + 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();
+}
+"""
+
+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>'
+
+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>'
+
+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">
+
+        <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>
+        </div>
+
+        <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>
+
+        <div class="app-main-row">
+            <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="#1E90FF" stroke-width="2" stroke-dasharray="4 3"/>
+                                <polygon points="12,62 30,40 42,50 54,34 68,62" fill="rgba(30,144,255,0.15)" stroke="#1E90FF" stroke-width="1.5"/>
+                                <circle cx="28" cy="30" r="6" fill="rgba(30,144,255,0.2)" stroke="#1E90FF" 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</div>
+                    <div class="examples-scroll">
+                        {EXAMPLE_CARDS_HTML}
+                    </div>
+                </div>
+            </div>
+
+            <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"><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>
+                        <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>
+
+        <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>
+            &middot; Open on <a href="https://github.com/PRITHIVSAKTHIUR/FireRed-Image-Edit-1.0-Fast" target="_blank">GitHub</a>
+        </div>
+
+        <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>
+    """)
+
+    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=30).launch(
+        css=css,
+        mcp_server=True,
+        ssr_mode=False,
+        show_error=True,
+        allowed_paths=["examples"],
+    )

pre-requirements.txt

@@ -0,0 +1 @@
+pip>=23.0.0

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,142 @@
+"""
+Paired with a good language model. Thanks!
+"""
+
+import torch
+from typing import Optional, Tuple
+from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
+
+try:
+    from kernels import get_kernel
+    _k = get_kernel("kernels-community/vllm-flash-attn3")
+    _flash_attn_func = _k.flash_attn_func
+except Exception as e:
+    _flash_attn_func = None
+    _kernels_err = e
+
+
+def _ensure_fa3_available():
+    if _flash_attn_func is None:
+        raise ImportError(
+            "FlashAttention-3 via Hugging Face `kernels` is required. "
+            "Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n"
+            f"{_kernels_err}"
+        )
+
+@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:
+    outputs, lse = _flash_attn_func(q, k, v, causal=causal)
+    return outputs
+
+@flash_attn_func.register_fake
+def _(q, k, v, **kwargs):
+    # two outputs:
+    # 1. output: (batch, seq_len, num_heads, head_dim)
+    # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32
+    meta_q = torch.empty_like(q).contiguous()
+    return meta_q #, q.new_empty((q.size(0), q.size(2), q.size(1)), dtype=torch.float32)
+
+
+class QwenDoubleStreamAttnProcessorFA3:
+    """
+    FA3-based attention processor for Qwen double-stream architecture.
+    Computes joint attention over concatenated [text, image] streams using vLLM FlashAttention-3
+    accessed via Hugging Face `kernels`.
+
+    Notes / limitations:
+    - General attention masks are not supported here (FA3 path). `is_causal=False` and no arbitrary mask.
+    - Optional windowed attention / sink tokens / softcap can be plumbed through if you use those features.
+    - Expects an available `apply_rotary_emb_qwen` in scope (same as your non-FA3 processor).
+    """
+
+    _attention_backend = "fa3"  # for parity with your other processors, not used internally
+
+    def __init__(self):
+        _ensure_fa3_available()
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        attn,  # Attention module with to_q/to_k/to_v/add_*_proj, norms, to_out, to_add_out, and .heads
+        hidden_states: torch.FloatTensor,                 # (B, S_img, D_model)  image stream
+        encoder_hidden_states: torch.FloatTensor = None,  # (B, S_txt, D_model)  text stream
+        encoder_hidden_states_mask: torch.FloatTensor = None,  # unused in FA3 path
+        attention_mask: Optional[torch.FloatTensor] = None,    # unused in FA3 path
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # (img_freqs, txt_freqs)
+    ) -> 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:
+            # FA3 kernel path here does not consume arbitrary masks; fail fast to avoid silent correctness issues.
+            raise NotImplementedError("attention_mask is not supported in this FA3 implementation.")
+
+        _ensure_fa3_available()
+
+        B, S_img, _ = hidden_states.shape
+        S_txt = encoder_hidden_states.shape[1]
+
+        # ---- QKV projections (image/sample stream) ----
+        img_q = attn.to_q(hidden_states)   # (B, S_img, D)
+        img_k = attn.to_k(hidden_states)
+        img_v = attn.to_v(hidden_states)
+
+        # ---- QKV projections (text/context stream) ----
+        txt_q = attn.add_q_proj(encoder_hidden_states)  # (B, S_txt, D)
+        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 (per your module contract) ----
+        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
+            # expects tensors shaped (B, S, H, D_h)
+            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 ----
+        # Shapes: (B, S_total, H, D_h)
+        q = torch.cat([txt_q, img_q], dim=1)
+        k = torch.cat([txt_k, img_k], dim=1)
+        v = torch.cat([txt_v, img_v], dim=1)
+
+        # FlashAttention-3 path expects (B, S, H, D_h) and returns (out, softmax_lse)
+        out = flash_attn_func(q, k, v, causal=False)  # out: (B, S_total, H, D_h)
+
+        # ---- Back to (B, S, D_model) ----
+        out = out.flatten(2, 3).to(q.dtype)
+
+        # Split back to 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,14 @@
+git+https://github.com/huggingface/accelerate.git
+git+https://github.com/huggingface/diffusers.git
+git+https://github.com/huggingface/peft.git
+transformers==4.57.6
+huggingface_hub
+sentencepiece
+torchvision
+kernels
+spaces
+hf_xet
+gradio
+torch
+numpy
+av