Add v0.2 adapters

.gitattributes

@@ -36,8 +36,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 assets/latency_t2i_1k_2k.png filter=lfs diff=lfs merge=lfs -text
 assets/download-1.png filter=lfs diff=lfs merge=lfs -text
 assets/edit_grid.jpg filter=lfs diff=lfs merge=lfs -text
-assets/Pruna-Q-Image-2.1-i2i-1024.png filter=lfs diff=lfs merge=lfs -text
-assets/Pruna-Q-Image-2.1-i2i-2048.png filter=lfs diff=lfs merge=lfs -text
-assets/Pruna-Q-Image-2.1-t2i-1024.png filter=lfs diff=lfs merge=lfs -text
-assets/Pruna-Q-Image-2.1-t2i-2048.png filter=lfs diff=lfs merge=lfs -text
-assets/Pruna-Q-Image-2.1.png filter=lfs diff=lfs merge=lfs -text
+assets/P-Q-Image-2.1-i2i-1024.png filter=lfs diff=lfs merge=lfs -text
+assets/P-Q-Image-2.1-i2i-2048.png filter=lfs diff=lfs merge=lfs -text
+assets/P-Q-Image-2.1-t2i-1024.png filter=lfs diff=lfs merge=lfs -text
+assets/P-Q-Image-2.1-t2i-2048.png filter=lfs diff=lfs merge=lfs -text
+assets/P-Q-Image-2.1.png filter=lfs diff=lfs merge=lfs -text
+assets/*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -16,7 +16,7 @@ tags:
 <!-- header start -->
 <!-- 200823 -->
 <p align="center">
-  <img src="assets/Pruna-Q-Image-2.1.png" width="800">
+  <img src="assets/P-Q-Image-2.1.png" width="800">
 </p>
 <!-- header end -->
 
@@ -43,12 +43,12 @@ tags:
 </p>
 <div align="center">
 
-<h1 style="color: #9334E9;">⚡ Pruna-Q-Image-2.1</h1>
+<h1 style="color: #9334E9;">⚡ P-Q-Image-2.1</h1>
 
 <h2>Few-step LoRA adapters for Qwen-Image-2.1</h2>
 
 <h3>
-  <span style="color: #9334E9;">5 or 8 steps</span>
+  <span style="color: #9334E9;">6 steps (v0.2), 5 or 8 steps (v0.1)</span>
   &nbsp;·&nbsp;
   <span style="color: #9334E9;">Up to 6.3x faster</span>
   &nbsp;·&nbsp;
@@ -59,29 +59,74 @@ tags:
 
 </div>
 
-**Pruna-Q-Image-2.1 is a set of LoRA adapters that let
+**P-Q-Image-2.1** is a set of LoRA adapters that let
 [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) generate
-and edit images in only 5 or 8 steps. The adapters load on top of the base
-model, so the pipeline, text encoder, and VAE stay unchanged. Training is based
-on DMD. Improved using Qwen.**
+and edit images in only 6 steps (v0.2), or 5 or 8 steps (v0.1). The adapters load on top of the base
+model, so the pipeline, text encoder, and VAE stay unchanged.
 
-> **v0.1: first release, work in progress.** Pruna-Q-Image-2.1 does not yet match the
-> visual quality of the base model. We are still improving the distillation and
-> will update this repository as new versions become available.
+> **v0.2: two new 6-step adapters.**
+>
+> - An advanced distillation for more faithful generation and editing.
+> - The first distillation removing need for prompt upsampling into the LoRA, thus significantly reducing end-to-end generation.
+> The v0.1 adapters remain available.
 
-## Variants
+<table>
+  <tr>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-t2i-1024-end-to-end.png" alt="Estimated end-to-end text-to-image latency at 1024 × 1024" width="100%">
+    </td>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-i2i-1024-end-to-end.png" alt="Estimated end-to-end editing latency at 1024 × 1024 with one input image" width="100%">
+    </td>
+  </tr>
+</table>
+
+## 6-step variants
+
+Pick **6-steps** for faithful base-model behavior, or **6-steps-upsampled** for
+improved visual quality from a short prompt, without a separate prompt-upsampling call.
+
+| Variant | File | Steps | Trade-off |
+|---|---|:---:|---|
+| `6-steps` | `p_q_image_2.1_6_steps_v0.2.safetensors` | 6 | **Advanced distillation.** Recommended for preserving inputs when editing. |
+| `6-steps-upsampled` | `p_q_image_2.1_6_steps_upsampled_v0.2.safetensors` | 6 | **Improved visual quality, shorter end-to-end generation.** The first distillation that removes the need for a separate prompt-upsampling step. |
+
+**6-steps** is an advanced 6-step distillation of the base model. It is intended
+as a drop-in replacement for the base model in existing generation and editing
+workflows, using the adapter's step count and sigma schedule. It improves on the
+v0.1 adapters in almost all categories and is particularly good at editing tasks
+where preserving the input matters.
+
+**6-steps-upsampled** distills the prompt upsampler recommended for
+Qwen-Image-2.1 into a 6-step LoRA. It improves visual quality without a separate
+prompt-upsampling call, which significantly reduces end-to-end generation time.
+It may not reproduce exactly what the actual prompt upsampler would do, and we
+will continue improving it.
+
+<details>
+<summary><b>See the 5-step vs 8-step adapters (v0.1)</b></summary>
+
+### v0.1
 
 Both adapters are v0.1. Pick one based on whether you need **quality** or **speed**.
 
-| File | Steps | Trade-off |
-|---|:---:|---|
-| `p_q_image_2.1_8step_v0.1.safetensors` | 8 | **Higher quality.** Recommended as the default. |
-| `p_q_image_2.1_5step_v0.1.safetensors` | 5 | **Higher speed**, with noticeably lower visual quality. |
+| Variant | File | Steps | Trade-off |
+|---|---|:---:|---|
+| `8-steps` | `p_q_image_2.1_8_steps_v0.1.safetensors` | 8 | **Higher quality.** Recommended as the default. |
+| `5-steps` | `p_q_image_2.1_5_steps_v0.1.safetensors` | 5 | **Higher speed**, with noticeably lower visual quality. |
 
 Each adapter is trained for its own sigma schedule. Load only one at a time.
+For the v0.2 adapters, do not reuse the v0.1 schedules.
+
+</details>
 
 ## Prompts, editing, and resolution
 
+For **6-steps**, use the same prompts and reference images you would use with the
+base model. For **6-steps-upsampled**, start with your original prompt; the adapter
+is intended to capture the benefit of prompt upsampling during generation.
+Detailed prompts still let you specify the subject, setting, lighting, and style.
+
 The adapters were trained at **1K resolution only**, with a mix of simple and
 upsampled prompts, text-to-image generation, and single- and multi-image editing
 with **up to 3 reference images**. Prompt upsampling is optional; detailed prompts
@@ -92,37 +137,6 @@ Higher resolutions (including 2K) and more reference images may work, but are
 outside the training coverage, and quality may vary. The 2K timings below show
 inference speed, not a guarantee of quality at that resolution.
 
-## Benchmarking
-
-Use a table; Hugging Face can otherwise render the `<img>` elements as separate blocks.
-
-<table>
-  <tr>
-    <td width="50%">
-      <img src="assets/Pruna-Q-Image-2.1-t2i-1024.png" width="100%">
-    </td>
-    <td width="50%">
-      <img src="assets/Pruna-Q-Image-2.1-t2i-2048.png" width="100%">
-    </td>
-  </tr>
-  <tr>
-    <td width="50%">
-      <img src="assets/Pruna-Q-Image-2.1-i2i-1024.png" width="100%">
-    </td>
-    <td width="50%">
-      <img src="assets/Pruna-Q-Image-2.1-i2i-2048.png" width="100%">
-    </td>
-  </tr>
-</table>
-
-Text-to-image with the official Qwen model-card example prompt, BF16, batch size 1,
-on one NVIDIA H100 80GB. Median of 3 requests after one warmup per configuration.
-Includes prompt encoding, denoising, and decoding; excludes PNG saving, model
-loading, and warmup. Base: 40 steps with KV cache on; Pruna-Q-Image-2.1: 5 or 8 steps with
-KV cache off, as configured for this benchmark. The examples below enable KV caching for
-Pruna-Q-Image-2.1; the chart has not been remeasured with that setting. LoRAs were unmerged;
-no CFG, compilation, or CPU offload. These timings do not imply equal image quality.
-
 ## Quickstart
 
 Runs text-to-image and image editing with the same pipeline. Needs a **CUDA GPU**.
@@ -132,6 +146,55 @@ pip install 'torch>=2.4.0' 'transformers>=5.17' accelerate peft pillow
 pip install git+https://github.com/huggingface/diffusers@6256aa7666cedd47443adc8f82da9a10e110b09c
 ```
 
+Use this setup for either v0.2 adapter.
+
+```python
+import torch
+from PIL import Image
+from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
+
+# Both v0.2 adapters use this unshifted six-step schedule.
+# The scheduler applies the resolution-dependent shift and appends sigma 0.
+STEPS = 6
+VARIANT = "6-steps"  # "6-steps" or "6-steps-upsampled"
+WEIGHT_NAME = {
+    "6-steps": "p_q_image_2.1_6_steps_v0.2.safetensors",
+    "6-steps-upsampled": "p_q_image_2.1_6_steps_upsampled_v0.2.safetensors",
+}[VARIANT]
+SIGMAS = [1.0, 0.925, 0.85, 0.625, 0.375, 0.125]
+
+pipe = QwenImage21Pipeline.from_pretrained(
+    "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
+).to("cuda")
+pipe.load_lora_weights(
+    "PrunaAI/P-Q-Image-2.1",
+    weight_name=WEIGHT_NAME,
+)
+
+# Shift the raw sigmas using the base model's resolution-dependent settings.
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+    pipe.scheduler.config,
+    use_dynamic_shifting=True,
+    shift_terminal=None,
+)
+```
+
+Both variants load on **`Qwen/Qwen-Image-2.1`** using **`QwenImage21Pipeline`**.
+Choose one LoRA with `VARIANT`; the `6-steps-upsampled` variant runs directly on your
+original prompt without a separate prompt-enhancer model. For a local export,
+replace `"PrunaAI/P-Q-Image-2.1"` with the folder containing the chosen
+safetensors file and keep `weight_name=WEIGHT_NAME`.
+
+Both files have **rank 64** across **232 linear targets**, including `proj_out`.
+Training used alpha 128; the exported B matrices include the `128 / 64 = 2`
+scale. Load at **LoRA strength 1.0** with the standard pipeline loader.
+
+
+<details>
+<summary><b>Check more info for the 5–8 step setup (v0.1)</b></summary>
+
+### v0.1
+
 ```python
 import torch
 from PIL import Image
@@ -149,8 +212,8 @@ pipe = QwenImage21Pipeline.from_pretrained(
     "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
 ).to("cuda")
 pipe.load_lora_weights(
-    "PrunaAI/Pruna-Q-Image-2.1",
-    weight_name=f"p_q_image_2.1_{STEPS}step_v0.1.safetensors",
+    "PrunaAI/P-Q-Image-2.1",
+    weight_name=f"p_q_image_2.1_{STEPS}_steps_v0.1.safetensors",
 )
 
 # Use the sigmas exactly as given: no extra shifting.
@@ -162,6 +225,89 @@ pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
 )
 ```
 
+</details>
+
+## Benchmarking
+
+
+<table>
+  <tr>
+    <td>
+      <img src="assets/P-Q-Image-2.1-v02-generation.png" alt="Six-step generation latency at 1024 and 2048 pixels" width="100%">
+    </td>
+  </tr>
+  <tr>
+    <td>
+      <img src="assets/P-Q-Image-2.1-prompt-upsampling.png" alt="Prompt-upsampling latency for text and one to three reference images" width="100%">
+    </td>
+  </tr>
+</table>
+
+One NVIDIA H100 80GB, BF16, batch size 1. v0.2 generation timings
+report the median of six measured requests pooled across both adapters (three
+per adapter), excluding one warmup per adapter and case.
+Generation includes prompt encoding, reference-image processing, denoising, and
+VAE decode. It excludes loading, warmup, fixture prep, saving, and network.
+No CFG, compilation, or CPU offload. **Diffusion KV caching is enabled**, except
+for the case marked with an asterisk below.
+
+\* **2048 × 2048 editing with 3 input images:** KV caching was disabled due to
+memory constraints on the H100 80GB. The 2048 multi-image cases use expandable
+CUDA memory segments (`PYTORCH_ALLOC_CONF=expandable_segments:True`).
+
+Prompt upsampling is timed separately with
+[`Qwen-Image-2.1-PE-T2I`](https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I)
+and [`Qwen-Image-2.1-PE-I2I`](https://huggingface.co/Qwen/Qwen-Image-2.1-PE-I2I).
+It covers request prep through parsed text, not image generation. vLLM runs
+eager, without prefix or image-processor caching, and with a fresh image ID so
+vision encoding is always counted. Autoregressive KV cache stays on.
+
+Text-to-image uses a fixed neon-sign prompt. Editing uses a fixed
+prompt and 1–3 reference images, square-cropped and resized. These cases are
+not matched to the v0.1 charts, and adding the two medians is not an
+end-to-end time.
+
+Generation: PyTorch 2.11.0, CUDA 12.8, Transformers 5.20.1, PEFT 0.21.2,
+Diffusers `6256aa7666cedd47443adc8f82da9a10e110b09c`. Upsampling: vLLM 0.19.1,
+PyTorch 2.10.0, Transformers 5.19.0. The v0.1 charts were not rerun here.
+Both enhancers use thinking with
+`temperature=1.0`, `top_p=0.95`, `top_k=20`, `min_p=0.0`. Text-to-image:
+`presence_penalty=1.5`, 16,256-token cap. Editing: `presence_penalty=0.0`,
+24,000-token cap. Reported only if all three rewrites complete and parse.
+
+<details>
+<summary><b>Check more info for the 5–8 step benchmarks (v0.1)</b></summary>
+
+### v0.1
+
+<table>
+  <tr>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-t2i-1024.png" width="100%">
+    </td>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-t2i-2048.png" width="100%">
+    </td>
+  </tr>
+  <tr>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-i2i-1024.png" width="100%">
+    </td>
+    <td width="50%">
+      <img src="assets/P-Q-Image-2.1-i2i-2048.png" width="100%">
+    </td>
+  </tr>
+</table>
+
+Text-to-image with the official Qwen model-card example prompt, BF16, batch size 1,
+on one NVIDIA H100 80GB. Median of 3 requests after one warmup per configuration.
+Includes prompt encoding, denoising, and decoding; excludes PNG saving, model
+loading, and warmup. Base: 40 steps with KV cache on; P-Q-Image-2.1: 5 or 8 steps with
+KV cache off, as configured for this benchmark. The examples below enable KV caching for
+P-Q-Image-2.1; the chart has not been remeasured with that setting. LoRAs were unmerged;
+no CFG, compilation, or CPU offload. These timings do not imply equal image quality.
+
+
 ## For Text-to-Image
 
 Longer, more descriptive prompts give better results.
@@ -199,6 +345,16 @@ image = pipe(
 
 ### Recommended settings
 
+- Use the unshifted six-step sigmas `[1.0, 0.925, 0.85, 0.625, 0.375, 0.125]` with dynamic shifting enabled
+and `shift_terminal=None`, as shown in the v0.2 setup.
+- The scheduler appends terminal zero.
+- Keep `true_cfg_scale=1.0` and LoRA strength 1.0.
+
+<details>
+<summary><b>Recommended 5–8 step settings (v0.1)</b></summary>
+
+### v0.1 settings
+
 - **Use the sigma schedule that matches your adapter**, and keep the scheduler
   at `shift=1.0` with dynamic shifting off so the sigmas are not shifted twice.
   - 8-step: `1 → 14/15 → 6/7 → 10/13 → 2/3 → 6/11 → 0.4 → 2/9 → 0`
@@ -210,8 +366,23 @@ image = pipe(
   the subject, setting, lighting, and style work noticeably better than short
   ones.
 
+</details>
+
 ## Limitations
 
+- The `6-steps-upsampled` adapter can improve visual quality without a separate
+  prompt-upsampling call, but may not exactly reproduce the output of the actual
+  prompt upsampler.
+- `6-steps` is intended to preserve the base model's behavior in fewer steps;
+  distillation does not guarantee identical outputs.
+- Use the step count and sigma schedule that match your v0.2 adapter. Other step
+  counts, schedules, or CFG values are not supported.
+
+<details>
+<summary><b>Known limits of the 5–8 step adapters (v0.1)</b></summary>
+
+### v0.1
+
 - This is a first version. Quality is below that of the base model.
 - The 5-step adapter is faster than the 8-step one, but its images are
   visibly worse.
@@ -219,6 +390,8 @@ image = pipe(
 - The adapters are only intended for the settings above. Other step counts,
   schedules, or CFG values are not supported.
 
+</details>
+
 ## What this is not
 
 - Not a standalone model. It needs the `Qwen/Qwen-Image-2.1` base weights.
@@ -229,7 +402,7 @@ image = pipe(
 
 The end distilled model is based on Qwen-Image-2.1.
 
-Pruna-Q-Image-2.1 is a derivative of Qwen-Image-2.1 and is distributed
+P-Q-Image-2.1 is a derivative of Qwen-Image-2.1 and is distributed
 under the **Qwen RESEARCH LICENSE AGREEMENT**. Review its use restrictions
 before using or redistributing the adapters.
 
@@ -238,7 +411,7 @@ Hangzhou Tongyi Laboratory Technology Co., Ltd. All Rights Reserved.
 
 ## What's next?
 
-- **Use Pruna-Q-Image-2.1 to generate and edit images with Qwen-Image-2.1 in a few steps.**
+- **Use P-Q-Image-2.1 to generate and edit images with Qwen-Image-2.1 in a few steps.**
 - Compress your own models with [Pruna](https://github.com/PrunaAI/pruna) and give us a ⭐️ for more efficiency!
 - Want to use our optimized image models right away?  Check [P-Image-Ideogram](https://www.pruna.ai/p-image-ideogram) and [P-Image-Edit](https://www.pruna.ai/p-image-edit).
 

p_q_image_2.1_6_steps_upsampled_v0.2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c86e69c1242625923d1454ba01ea6cef900f80ee891c5fa245173ed4abac9c2
+size 688027008

p_q_image_2.1_6_steps_v0.2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d10a0ab365df3829d460be9196f6a8a7d0e8f5498f5a3f25d924b8f9ee8c94bb
+size 688027008