Add v0.2 adapters

.gitattributes

@@ -41,3 +41,5 @@ 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/*.png filter=lfs diff=lfs merge=lfs -text
+assets/Pruna-Qwen-Image-2.1.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -48,7 +48,7 @@ tags:
 <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;
@@ -61,16 +61,133 @@ tags:
 
 **Pruna-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
+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. Training is based
-on DMD. Improved using Qwen.**
+on DMD.**
+
+> **v0.2: two new 6-step adapters.** A basic distillation for faithful generation
+> and editing, and a first attempt at distilling prompt upsampling into the LoRA.
+> The v0.1 adapters remain available.
 
 > **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.
 
+## Quickstart
+
+Runs text-to-image and image editing with the same pipeline. Needs a **CUDA GPU**.
+
+```bash
+pip install 'torch>=2.4.0' 'transformers>=5.17' accelerate peft pillow
+pip install git+https://github.com/huggingface/diffusers@6256aa7666cedd47443adc8f82da9a10e110b09c
+```
+
+### v0.2
+
+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 = "basic"  # "basic" or "upsampled"
+WEIGHT_NAME = {
+    "basic": "p_q_image_2.1_basic_6_steps_v0.2.safetensors",
+    "upsampled": "p_q_image_2.1_upsampled_6_steps_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/Pruna-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 prompt-upsampled variant runs directly on your
+original prompt without a separate prompt-enhancer model. For a local export,
+replace `"PrunaAI/Pruna-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.
+
+For v0.2 benchmark settings, pass `output_resolution=1024` and
+`use_kv_cache=False` in the generation and editing calls below.
+
+### v0.1
+
+```python
+import torch
+from PIL import Image
+from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
+
+STEPS = 8  # 8 for higher quality, 5 for higher speed
+
+# The terminal sigma 0 is appended by the scheduler.
+SIGMAS = {
+    5: [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
+    8: [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
+}[STEPS]
+
+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",
+)
+
+# Use the sigmas exactly as given: no extra shifting.
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+    pipe.scheduler.config,
+    use_dynamic_shifting=False,
+    shift=1.0,
+    shift_terminal=None,
+)
+```
+
 ## Variants
 
+### v0.2
+
+Pick **basic** for faithful base-model behavior, or **prompt-upsampled** for
+improved visual quality from a short prompt.
+
+| File | Steps | Trade-off |
+|---|:---:|---|
+| `p_q_image_2.1_basic_6_steps_v0.2.safetensors` | 6 | **Pure few-step distillation.** Recommended for preserving inputs when editing. |
+| `p_q_image_2.1_upsampled_6_steps_v0.2.safetensors` | 6 | **Improved visual quality.** A first attempt at distilling prompt upsampling into the adapter. |
+
+**Basic** is a pure 6-step distillation of the base model, without any additional
+post-training. 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.
+
+**Prompt-upsampled** attempts to distill the prompt upsampler recommended for
+Qwen-Image-2.1 into a 6-step LoRA. It improves visual quality without a separate
+prompt-upsampling call. This is our first attempt: it may not reproduce exactly
+what the actual prompt upsampler would do, and we will continue improving it.
+
+### v0.1
+
 Both adapters are v0.1. Pick one based on whether you need **quality** or **speed**.
 
 | File | Steps | Trade-off |
@@ -79,9 +196,17 @@ Both adapters are v0.1. Pick one based on whether you need **quality** or **spee
 | `p_q_image_2.1_5step_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.
 
 ## Prompts, editing, and resolution
 
+For **v0.2**:
+
+For **basic**, use the same prompts and reference images you would use with the
+base model. For **prompt-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
@@ -94,7 +219,65 @@ 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.
+### v0.2: 6-step LoRA and prompt upsampling
+
+<table>
+  <tr>
+    <td>
+      <img src="assets/Pruna-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/Pruna-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 per worker, BF16, batch size 1. Median of 3 requests after
+one warmup per configuration, with seed 42. Wall-clock time is measured around
+the complete request, with CUDA synchronization before and after. Generation
+includes prompt encoding, reference-image
+processing, denoising, and VAE decoding to a PIL image. It excludes model and
+adapter loading, warmup, fixture preparation, file saving, and network transport.
+No CFG, compilation, CPU offload, or diffusion KV caching.
+
+Prompt upsampling is measured separately using
+[`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),
+with their task-specific sampling settings and thinking enabled. Timing includes
+request preparation from in-memory images, image and text processing, text
+generation, and response parsing. It does not include image generation. The
+upsampler uses vLLM in eager mode, without prefix or image-processor caching.
+Each request receives fresh image identifiers, so vision encoding is included
+on repeated requests too. Autoregressive KV caching remains enabled; it is
+separate from diffusion KV caching.
+
+Each editing case uses a fixed original prompt and 1, 2, or 3 reference images,
+square-cropped and resized to the requested resolution before timing.
+Text-to-image uses the official Qwen model-card example prompt. Latency depends
+on the prompt, image
+content, and generated rewrite length. These cases are not a controlled
+comparison with the v0.1 chart inputs. Adding the two medians does not
+measure an end-to-end prompt-upsampling-plus-generation request.
+
+Software and prompt-upsampling settings
+
+Generation uses PyTorch 2.11.0 with CUDA 12.8, Transformers 5.20.1, PEFT 0.21.2,
+and Diffusers commit `6256aa7666cedd47443adc8f82da9a10e110b09c`, the same Diffusers
+revision used in the quickstart. Prompt upsampling uses vLLM 0.19.1, PyTorch
+2.10.0, and Transformers 5.19.0. These runtime versions are recorded for the new
+measurements; the v0.1 charts have not been rerun in this environment.
+
+Both prompt enhancers use their checkpoint's system prompt, thinking enabled,
+`temperature=1.0`, `top_p=0.95`, `top_k=20`, and `min_p=0.0`. Text-to-image uses
+`presence_penalty=1.5` and a 16,256-token output limit; editing uses
+`presence_penalty=0.0` and a 24,000-token output limit. A configuration is
+reported only if all three measured rewrites complete and parse successfully.
+No requests are silently retried.
+
+### v0.1
+
 
 <table>
   <tr>
@@ -123,45 +306,6 @@ KV cache off, as configured for this benchmark. The examples below enable KV cac
 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**.
-
-```bash
-pip install 'torch>=2.4.0' 'transformers>=5.17' accelerate peft pillow
-pip install git+https://github.com/huggingface/diffusers@6256aa7666cedd47443adc8f82da9a10e110b09c
-```
-
-```python
-import torch
-from PIL import Image
-from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
-
-STEPS = 8  # 8 for higher quality, 5 for higher speed
-
-# The terminal sigma 0 is appended by the scheduler.
-SIGMAS = {
-    5: [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
-    8: [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
-}[STEPS]
-
-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",
-)
-
-# Use the sigmas exactly as given: no extra shifting.
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
-    pipe.scheduler.config,
-    use_dynamic_shifting=False,
-    shift=1.0,
-    shift_terminal=None,
-)
-```
-
 ## For Text-to-Image
 
 Longer, more descriptive prompts give better results.
@@ -199,6 +343,13 @@ image = pipe(
 
 ### Recommended settings
 
+For **v0.2**, 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.
+
+The **v0.1** settings are:
+
 - **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`
@@ -212,6 +363,17 @@ image = pipe(
 
 ## Limitations
 
+### v0.2
+
+- The prompt-upsampled adapter is a first attempt. It can improve visual quality,
+  but may not exactly reproduce the output of the actual prompt upsampler.
+- Basic 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.
+
+### 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.

p_q_image_2.1_basic_6_steps_v0.2.safetensors

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

p_q_image_2.1_upsampled_6_steps_v0.2.safetensors

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