nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[new-model] Add initial Qwen-Image-2.1 inference (#1918)

Browse Source 
Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com>
This commit is contained in:
Loay Rashid
2026-10-25 17:28:49 +00:00
committed by GitHub
co-authored by Aryan Kumar
parent a92c5c1578
commit 0d05b09c57
44 changed files with 4746 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/scripts/plan_merge_ci.py
+5
View File
@@ -131,6 +131,11 @@ class FamilyCoverage:
FAMILY_COVERAGE = (
FamilyCoverage(
re.compile(r"(^|[/_.-])qwen[_-]?image[_-]?2[._-]?1([/_.-]|$)"),
("test_qwen_image21.py", ),
("test_qwen_image21_similarity.py", ),
),
FamilyCoverage(
re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
("test_flux2_klein.py", ),
docs-site/lib/cookbook-pages.ts
+3
View File
@@ -111,6 +111,7 @@ const familyInfo: Record<string, { title: string; path: string; logo: string; ca
kandinsky5: { title: "Kandinsky 5", path: "kandinsky5", logo: "kandinsky", category: "Video" },
kandinsky6: { title: "Kandinsky 6", path: "kandinsky6", logo: "kandinsky", category: "Video" },
flux: { title: "FLUX", path: "flux", logo: "black-forest-labs", category: "Image" },
qwen_image21: { title: "Qwen-Image-2.1", path: "qwen-image21", logo: "qwen", category: "Image" },
sd35: { title: "Stable Diffusion", path: "stable-diffusion", logo: "stabilityai", category: "Image" },
stable_audio: { title: "Stable Audio", path: "stable-audio", logo: "stabilityai", category: "Audio" },
mmaudio: { title: "MMAudio", path: "mmaudio", logo: "fastvideo", category: "Audio" },
@@ -132,6 +133,7 @@ const TASK_GROUPS: Array<[RegExp, string]> = [
[/video\/text to audio|video.*to audio/, "Video to audio"],
[/text to audio/, "Text to audio"],
[/text to image/, "Text to image"],
[/image to image/, "Image to image"],
[/reference media to video/, "Reference to video"],
[/first\/last frame to video/, "Image to video"],
[/text or image to video/, "Text/image to video"],
@@ -161,6 +163,7 @@ export const modelDirectory = [...new Set(sourceRecipes.map((recipe) => recipe.m
"tencent-hunyuan": "Tencent",
kandinsky: "Kandinsky Lab",
"black-forest-labs": "Black Forest Labs",
qwen: "Qwen",
stabilityai: "Stability AI",
nvidia: "NVIDIA",
};
docs-site/navigation.json
+4
View File
@@ -253,6 +253,10 @@
"label": "FLUX",
"href": "/cookbook/flux"
},
{
"label": "Qwen-Image-2.1",
"href": "/cookbook/qwen-image21"
},
{
"label": "Stable Diffusion",
"href": "/cookbook/stable-diffusion"
docs-site/tests/cookbook.spec.ts
+25 -5
View File
@@ -223,7 +223,7 @@ test("inherited values, explicit false, zero and null retain their meanings", as
test("catalog discovery opens dedicated model recipes and filters by family and task", async ({ page }) => {
await page.goto("/FastVideo/cookbook");
const cards = page.locator(".recipe-card");
await expect(cards).toHaveCount(36);
await expect(cards).toHaveCount(37);
await expect(cards.first()).toContainText("FastH3 OmniRef");
await page.getByLabel("Model family", { exact: true }).selectOption("Wan");
await page.getByLabel("Search recipes").fill("Wan2.1 I2V");
@@ -277,7 +277,7 @@ test("output tabs support keyboard navigation and precision chips use the backen
await page.getByLabel("Search recipes").fill("does not exist");
await expect(page.getByRole("heading", { name: "No matching recipes" })).toBeVisible();
await page.getByRole("button", { name: "Clear filters" }).click();
await expect(page.locator(".recipe-card")).toHaveCount(36);
await expect(page.locator(".recipe-card")).toHaveCount(37);
});
test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => {
@@ -297,7 +297,7 @@ test("cookbook cards and controls remain readable in dark mode and reflow at 320
test("image and audio models have logos and dedicated source recipe pages", async ({ page }) => {
await page.goto("/FastVideo/cookbook");
await page.getByRole("button", { name: "Image", exact: true }).click();
await expect(page.locator(".recipe-card")).toHaveCount(4);
await expect(page.locator(".recipe-card")).toHaveCount(5);
const card = page.locator(".recipe-card").filter({ hasText: "FLUX.2 Klein" });
await expect(card.locator("img")).toBeVisible();
await card.click();
@@ -309,6 +309,26 @@ test("image and audio models have logos and dedicated source recipe pages", asyn
await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
await page.getByLabel("Seed", { exact: true }).fill("0");
await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
await page.goto("/FastVideo/cookbook/qwen-image21");
const qwen = page.locator(".recipe-card").filter({ hasText: "Qwen-Image-2.1" });
await expect(qwen.locator("img")).toHaveAttribute("src", /\/qwen\.webp$/);
await qwen.click();
await expect(page.getByRole("heading", { level: 1 })).toHaveText("Qwen-Image-2.1");
await expect(page.getByRole("button", { name: "Serving", exact: true })).toHaveCount(0);
await page.getByLabel("Height", { exact: true }).fill("512");
await page.getByRole("tab", { name: "config.yaml", exact: true }).click();
const qwenConfig = page.getByLabel("Source example", { exact: true });
await expect(qwenConfig).toContainText("model_path: Qwen/Qwen-Image-2.1");
await expect(qwenConfig).toContainText("height: 512");
await expect(qwenConfig).toContainText("dit_layerwise: true");
await expect(page.getByRole("button", { name: "Download YAML", exact: true })).toBeEnabled();
await page
.getByRole("group", { name: "Recipe variant", exact: true })
.getByRole("button", { name: "Qwen-Image-2.1 references and editing", exact: true })
.click();
await page.getByRole("tab", { name: "Commands", exact: true }).click();
await expect(page.getByLabel("Inference command", { exact: true })).toContainText("--reference source.png");
await expect(page.getByRole("button", { name: "Copy example + command", exact: true })).toBeEnabled();
await page.goto("/FastVideo/cookbook");
await page.getByRole("button", { name: "Audio", exact: true }).click();
await expect(page.locator(".recipe-card")).toHaveCount(3);
@@ -584,11 +604,11 @@ test("directory filters round-trip through the URL", async ({ page }) => {
await page.goto("/FastVideo/cookbook?category=Image&runtime=cuda");
await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
await expect(page.locator(".recipe-card")).toHaveCount(4);
await expect(page.locator(".recipe-card")).toHaveCount(5);
await page.reload();
await expect(page.getByRole("button", { name: "Image", exact: true })).toHaveAttribute("aria-pressed", "true");
await expect(page.getByLabel("Runtime", { exact: true })).toHaveValue("cuda");
await expect(page.locator(".recipe-card")).toHaveCount(4);
await expect(page.locator(".recipe-card")).toHaveCount(5);
});
test("a saved target missing from the catalog stays visible and clearable", async ({ page }) => {
docs/assets/cookbook-recipes.json
+39 -1
View File
@@ -1,5 +1,5 @@
{
"version": 18,
"version": 19,
"recipes": [
{
"id": "fastwan21-t2v",
@@ -429,6 +429,44 @@
"expected_artifact": "PNG images under outputs/flux_dev/samples/",
"limitations": ["FLUX.1 is loadable by ID but registers no model_family in fastvideo/registry.py; it is grouped under FLUX for documentation only."]
},
{
"id": "qwen-image21-t2i",
"family": "qwen_image21",
"stage": "inference",
"task": "Text to image",
"label": "Qwen-Image-2.1",
"summary": "Generate a 1024 × 1024 image in 40 steps on one CUDA GPU, with the CPU offload settings recorded in the maintained YAML.",
"model": "Qwen/Qwen-Image-2.1",
"source": "examples/inference/basic/qwen_image21_t2i.yaml",
"command": "fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml",
"gpu_types": ["NVIDIA"],
"hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "RGBA PNG image at outputs/qwen_image21/teapot.png",
"limitations": ["The source uses DiT layerwise offload, text-encoder and VAE offload, and lazy loading. GPU memory requirements have not been established for this recipe."],
"related": ["qwen-image21-i2i"]
},
{
"id": "qwen-image21-i2i",
"family": "qwen_image21",
"stage": "inference",
"task": "Image to image",
"label": "Qwen-Image-2.1 references and editing",
"summary": "Generate or edit an image from a prompt and up to ten ordered reference images, using the maintained single-GPU Python example.",
"model": "Qwen/Qwen-Image-2.1",
"source": "examples/inference/basic/basic_qwen_image21.py",
"command": "python examples/inference/basic/basic_qwen_image21.py --reference source.png --prompt 'Change the shirt to blue; retain the subject and composition.' --output outputs/qwen_image21/edit.png",
"gpu_types": ["NVIDIA"],
"hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"},
"evidence": "Source-backed",
"expected_artifact": "RGBA PNG image at outputs/qwen_image21/edit.png",
"limitations": [
"Replace source.png with an existing image. Repeat --reference for additional images and describe their order in the prompt.",
"Mask images use the same ordered-reference path; they do not impose hard pixel-preservation constraints.",
"The source enables CPU offload and lazy loading. This entry does not establish a minimum GPU memory requirement."
],
"related": ["qwen-image21-t2i"]
},
{
"id": "sd35-medium-t2i",
"family": "sd35",
docs/assets/logos/SOURCES.md
+1
View File
@@ -13,6 +13,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo.
| `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header |
| `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers |
| `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header |
| `qwen.webp` | [Official Qwen Hugging Face organization avatar](https://huggingface.co/Qwen) | Qwen-Image-2.1 family card and page header |
| `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header |
| `stabilityai.webp` | [Official Stability AI Hugging Face organization avatar](https://huggingface.co/stabilityai) | Stable Diffusion and Stable Audio cards and page headers |
| `fastvideo.webp` | [Official FastVideo Hugging Face organization avatar](https://huggingface.co/FastVideo) | Matrix Game and MMAudio cards (converted weights published by this org) |
docs/assets/logos/qwen.webp
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.8 KiB

docs/cookbook/index.md
+1
View File
@@ -7,6 +7,7 @@ The interactive catalog groups models by family and generates YAML for supported
- [Wan and FastMetal](wan.md)
- [LTX](ltx.md)
- [FLUX](flux.md)
- [Qwen-Image-2.1](qwen-image21.md)
- [Stable Diffusion](stable-diffusion.md)
- [Stable Audio](stable-audio.md)
- [MMAudio](mmaudio.md)
docs/cookbook/qwen-image21.md
+108
View File
@@ -0,0 +1,108 @@
# Qwen-Image-2.1 recipes
<RecipeFamily family="qwen_image21" />
Generate images, edit an image, or combine up to ten ordered references with
`Qwen/Qwen-Image-2.1`. Outputs are RGBA PNGs. Request a transparent background
in the prompt when needed.
## Run on a GPU
Follow the [CUDA installation guide](../getting_started/installation/gpu.md), then
run from the repository root on one CUDA GPU (`tp_size=sp_size=1`). The examples
use CPU offload and VAE tiling. The loader downloads the complete checkpoint;
use `--model-path /path/to/Qwen-Image-2.1` for local weights. No conversion is needed.
```bash
python examples/inference/basic/basic_qwen_image21.py \
--prompt 'A red ceramic teapot, isolated on a transparent background.' \
--output outputs/teapot.png
```
This saves a 1024×1024 PNG after 40 denoising steps. For instruction editing,
replace `source.png` with your image:
```bash
python examples/inference/basic/basic_qwen_image21.py \
--reference source.png \
--prompt 'Change the shirt to blue; retain the subject and composition.' \
--output outputs/edit.png
```
Repeat `--reference` for reference-to-image generation, annotations or masks:
```bash
python examples/inference/basic/basic_qwen_image21.py \
--reference subject.png --reference style.png \
--prompt 'Paint the subject from image 1 in the style of image 2.'
```
A mask is another reference interpreted through the prompt; pixels outside it
are not guaranteed to remain identical.
## Configuration
- Set `--height`, `--width`, `--steps` and `--seed` as needed. Dimensions must be
positive multiples of 32; use `--height 2048 --width 2048` for 2K output.
- `--reference-resolution` defaults to 1024 and must be a positive multiple of
32. References retain their aspect ratio and are resized to approximately
that resolution squared in pixels, independently of output dimensions.
- Prefix KV caching avoids recomputing text/reference tokens each step. The
Python example caches on CPU; `--no-use-kv-cache` disables reuse.
- For two-branch CFG, set `--true-cfg-scale` above 1 and supply
`--negative-prompt` (an empty string is valid). Embedded guidance is not used.
For the CLI, use the checked-in [T2I YAML](../../examples/inference/basic/qwen_image21_t2i.yaml)
or [editing YAML](../../examples/inference/basic/qwen_image21_edit.yaml) with
`fastvideo generate --config <file>`. The T2I YAML keeps prefix KV caches on CUDA;
the Python example and editing YAML default to CPU caches. See
[Configuration](../inference/configuration.md) for the Python API and YAML options.
## Attention
Dense image attention uses FastVideo's shared automatic backend selector, which
prefers FlashAttention on supported CUDA hardware. Masked text attention uses
SDPA. To request FA4 with its dependencies installed on supported hardware:
```bash
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
export FASTVIDEO_FA4=1
```
Use `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` and `FASTVIDEO_FA4=0` for upstream
parity comparisons. Different attention kernels can produce different pixels.
See [Optimizations](../inference/optimizations.md) for backend and offload settings.
## Measured memory
Peaks below are in **GiB**, measured on one **NVIDIA GB200** with BF16,
1024×1024 output, seed 42, compilation disabled and two 40-step requests per
process. Image cases use one reference. These measurements do not establish
minimum hardware capacity; an actual 24 GB GPU has not been verified.
| Configuration | KV cache | GPU allocated | GPU reserved | Device VRAM including driver | Host process-tree RSS |
| --- | --- | ---: | ---: | ---: | ---: |
| T2I YAML, SDPA, offloaded, tiled | CUDA | 13.28 | 22.54 | 24.19 | 43.87 |
| Plain I2I Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 48.29 |
| Editing Python, SDPA, offloaded, tiled | CPU | 13.29 | 16.11 | 17.76 | 49.25 |
| T2I, FA4, resident models, untiled VAE | CUDA | 32.77 | 35.50 | 37.09 | 13.12 |
Measurements cover loading, cold/warm requests, PNG export and shutdown.
Allocated/reserved values use PyTorch counters; device VRAM and simultaneous
process-tree RSS were sampled every 20 ms. Device VRAM includes the GB200's
0.72 GiB driver reservation. Allow host RAM headroom beyond process RSS.
Offloaded rows explicitly use SDPA; the resident FA4 row disables offload and
tiling and is a separate configuration, not a controlled latency comparison.
## Validation and scope
On GB200, 40-step upstream SDPA parity passed for eleven 1024×1024 cases:
T2I, plain I2I, editing, reference-to-image, RGBA generation/editing, annotations,
mask images, true CFG, uncached execution and ten references. Separate 2048×2048
T2I also passed. Denoising trajectories and output pixels matched exactly.
FA4 block and T2I/edit image regressions passed against their matching references.
See the [validation runbook](../../tests/local_tests/qwen_image21/README.md)
for reproducible commands, pinned revisions and numerical diagnostics.
The supported scope is single GPU, batch one and unquantized inference.
SP/TP, LoRA, prompt rewriting and HTTP serving are not included.
docs/design/inference_schema_parity_inventory.yaml
+10
View File
@@ -149,6 +149,7 @@ surfaces:
preprocess_text_funcs: "Internal text preprocessing hooks."
postprocess_text_funcs: "Internal text postprocessing hooks."
scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
output_channels: "Decoded pixel channel count fixed by the model; used to allocate the output buffer."
pipeline_config_extensions:
preset_owned:
@@ -442,6 +443,15 @@ surfaces:
sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
sample_fps:
sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig]
reference_resolution:
target: generator.pipeline.experimental.reference_resolution
sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
use_kv_cache:
target: generator.pipeline.experimental.use_kv_cache
sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
kv_cache_device:
target: generator.pipeline.experimental.kv_cache_device
sources: [fastvideo.configs.pipelines.qwen_image21.QwenImage21PipelineConfig]
compatibility_only:
batch_size: "Gen3C inference-only tuning field pending typed batching design."
gradient_checkpointing: "Gen3C inference-only compatibility field pending typed batching design."
docs/generate_examples.py
+1
View File
@@ -36,6 +36,7 @@ COOKBOOK_FAMILIES = {
"kandinsky5",
"kandinsky6",
"flux",
"qwen_image21",
"sd35",
"minimax_h3",
"stable_audio",
docs/inference/support_matrix.md
+1
View File
@@ -69,6 +69,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
| matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
| minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br />[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br />[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
| sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) |
| qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) |
| stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) |
| stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) |
| wan | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | T2V | [basic.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic.py) |
examples/inference/basic/basic_qwen_image21.py
+73
View File
@@ -0,0 +1,73 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 generation, editing, references and transparent PNG output."""
import argparse
from fastvideo import VideoGenerator
from fastvideo.api import (
EngineConfig,
GenerationRequest,
GeneratorConfig,
InputConfig,
OffloadConfig,
OutputConfig,
ParallelismConfig,
PipelineSelection,
SamplingConfig,
)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-path", default="Qwen/Qwen-Image-2.1")
parser.add_argument("--revision")
parser.add_argument("--prompt", default="A red ceramic teapot, isolated on a transparent background.")
parser.add_argument("--reference", action="append", default=[], help="Ordered reference image; repeat up to ten times")
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--true-cfg-scale", type=float, default=1.0)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--reference-resolution", type=int, default=1024)
parser.add_argument("--steps", type=int, default=40)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--output", default="outputs/qwen_image21/image.png")
args = parser.parse_args()
if len(args.reference) > 10:
parser.error("At most ten ordered reference images are supported")
workload = "i2i" if args.reference else "t2i"
generator = VideoGenerator.from_config(
GeneratorConfig(
model_path=args.model_path,
revision=args.revision,
engine=EngineConfig(
num_gpus=1,
parallelism=ParallelismConfig(tp_size=1, sp_size=1),
offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True, vae=True,
lazy_module_load=True),
),
pipeline=PipelineSelection(
workload_type=workload,
experimental={
"reference_resolution": args.reference_resolution,
"use_kv_cache": args.use_kv_cache,
},
),
))
try:
generator.generate(
GenerationRequest(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
inputs=InputConfig(references=args.reference or None),
sampling=SamplingConfig(height=args.height, width=args.width, num_frames=1, fps=1,
num_inference_steps=args.steps, guidance_scale=1.0,
true_cfg_scale=args.true_cfg_scale, seed=args.seed),
output=OutputConfig(output_path=args.output, save_video=True, return_frames=False),
))
finally:
generator.shutdown()
if __name__ == "__main__":
main()
examples/inference/basic/qwen_image21_edit.yaml
+39
View File
@@ -0,0 +1,39 @@
# Replace source.png and mask.png with your own images. The mask is an ordered
# reference interpreted by the model, not a hard pixel-preservation constraint.
generator:
model_path: Qwen/Qwen-Image-2.1
engine:
num_gpus: 1
parallelism:
tp_size: 1
sp_size: 1
offload:
dit: true
dit_layerwise: true
text_encoder: true
vae: true
lazy_module_load: true
pipeline:
workload_type: i2i
experimental:
reference_resolution: 1024
use_kv_cache: true
request:
prompt: In image 1, turn the object highlighted by the mask in image 2 blue.
inputs:
references:
- source.png
- mask.png
sampling:
height: 1024
width: 1024
num_frames: 1
fps: 1
num_inference_steps: 40
guidance_scale: 1.0
true_cfg_scale: 1.0
seed: 42
output:
output_path: outputs/qwen_image21/edit.png
save_video: true
return_frames: false
examples/inference/basic/qwen_image21_t2i.yaml
+40
View File
@@ -0,0 +1,40 @@
# Run on an allocated CUDA GPU after installing FastVideo and providing weights.
# fastvideo generate --config examples/inference/basic/qwen_image21_t2i.yaml
generator:
model_path: Qwen/Qwen-Image-2.1
engine:
num_gpus: 1
compile:
enabled: false
backend: inductor
fullgraph: false
parallelism:
tp_size: 1
sp_size: 1
offload:
dit: true
dit_layerwise: true
text_encoder: true
vae: true
lazy_module_load: true
pipeline:
workload_type: t2i
experimental:
kv_cache_device: cuda
reference_resolution: 1024
use_kv_cache: true
request:
prompt: A red ceramic teapot, isolated on a transparent background.
sampling:
height: 1024
width: 1024
num_frames: 1
fps: 1
num_inference_steps: 40
guidance_scale: 1.0
true_cfg_scale: 1.0
seed: 42
output:
output_path: outputs/qwen_image21/teapot.png
save_video: true
return_frames: false
fastvideo/configs/models/dits/__init__.py
+2 -1
View File
@@ -18,6 +18,7 @@ from fastvideo.configs.models.dits.mmaudio import (
)
from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
from fastvideo.models.wan.config import WanVideoConfig
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
from fastvideo.configs.models.dits.hyworld import HYWorldConfig
from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig
@@ -33,5 +34,5 @@ __all__ = [
"StableAudioConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotWorldFastVideoConfig",
"LingBotVideoConfig", "MiniMaxH3Config", "MMAudioArchConfig", "MMAudioTransformerConfig",
"MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS", "MMAUDIO_44K_TRAINING_VARIANTS",
"get_mmaudio_transformer_config"
"get_mmaudio_transformer_config", "QwenImage21ArchConfig", "QwenImage21Config"
]
fastvideo/configs/models/dits/qwen_image21.py
+55
View File
@@ -0,0 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
"""Architecture and checkpoint loading configuration for Qwen-Image-2.1."""
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.platforms import AttentionBackendEnum
def _is_transformer_block(name: str, module) -> bool:
parts = name.split(".")
return len(parts) >= 2 and parts[-2] == "transformer_blocks" and parts[-1].isdigit()
@dataclass
class QwenImage21ArchConfig(DiTArchConfig):
patch_size: int = 1
in_channels: int = 64
out_channels: int = 64
num_layers: int = 32
attention_head_dim: int = 128
num_attention_heads: int = 32
context_in_dim: int = 4096
mlp_ratio: int = 3
axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
eps: float = 1e-6
causal_condition: bool = True
cast_prompt_embeds_to_dit_dtype: bool = True
_fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_block])
_supported_attention_backends: tuple[AttentionBackendEnum,
...] = (AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.SAGE_ATTN)
param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
reverse_param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
def __post_init__(self) -> None:
axes_dims_rope = tuple(self.axes_dims_rope)
if len(axes_dims_rope) != 3 or any(dim <= 0 or dim % 2 for dim in axes_dims_rope):
raise ValueError("axes_dims_rope must contain three positive even dimensions")
self.axes_dims_rope = (axes_dims_rope[0], axes_dims_rope[1], axes_dims_rope[2])
if sum(self.axes_dims_rope) != self.attention_head_dim:
raise ValueError("axes_dims_rope must sum to attention_head_dim")
if self.num_layers <= 0 or self.num_attention_heads <= 0:
raise ValueError("num_layers and num_attention_heads must be positive")
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.out_channels
@dataclass
class QwenImage21Config(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=QwenImage21ArchConfig)
prefix: str = "QwenImage21"
fastvideo/configs/models/encoders/__init__.py
+3
View File
@@ -15,6 +15,7 @@ from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig)
from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLArchConfig, QwenImage21Qwen3VLConfig
from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
from fastvideo.configs.models.encoders.stable_audio_conditioner import (
StableAudioConditionerArchConfig,
@@ -49,6 +50,8 @@ __all__ = [
"StableAudioConditionerConfig",
"T5GemmaEncoderConfig",
"Qwen3TextConfig",
"QwenImage21Qwen3VLArchConfig",
"QwenImage21Qwen3VLConfig",
"Mistral3TextConfig",
"LingBotWorld2UMT5ArchConfig",
"LingBotWorld2UMT5Config",
fastvideo/configs/models/encoders/qwen_image21.py
+79
View File
@@ -0,0 +1,79 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen3-VL architecture used for Qwen-Image-2.1 conditioning."""
from dataclasses import dataclass, field
from typing import Any
from fastvideo.configs.models.base import ModelConfig
from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (
MiniMaxH3Qwen3VLArchConfig,
MiniMaxH3Qwen3VLConfig,
_OFFICIAL_ARCHITECTURES,
_VISION_CONFIG_MAPPING,
)
@dataclass
class QwenImage21Qwen3VLArchConfig(MiniMaxH3Qwen3VLArchConfig):
architectures: list[str] = field(default_factory=lambda: ["QwenImage21Qwen3VLConditioner"])
hidden_size: int = 4096
intermediate_size: int = 12288
num_hidden_layers: int = 36
output_hidden_state_index: int = 36
num_hidden_layers_override: int | None = None
num_attention_heads: int = 32
vision_out_hidden_size: int = 4096
# The image pipeline's processor owns prompt lengths, including reference
# image tokens; the H3 text-only truncation limit is inappropriate here.
text_len: int = 262144
require_processor: bool = True
def __post_init__(self) -> None:
super().__post_init__()
if self.output_hidden_state_index != self.num_hidden_layers:
raise ValueError("Qwen-Image-2.1 conditioning requires the last decoder layer before final normalization")
if self.num_hidden_layers_override is not None:
raise ValueError("Qwen-Image-2.1 conditioning requires all text decoder layers")
if self.hidden_act != "silu" or self.vision_hidden_act != "gelu_pytorch_tanh":
raise ValueError("Qwen-Image-2.1 requires SiLU text MLPs and tanh-approximate GELU vision MLPs")
if (self.rope_scaling or {}).get("rope_type", "default") != "default":
raise ValueError("Qwen-Image-2.1 conditioning supports default multimodal RoPE")
self.tokenizer_kwargs = {
"padding": True,
"padding_side": "left",
"truncation": False,
"return_tensors": "pt",
}
@dataclass
class QwenImage21Qwen3VLConfig(MiniMaxH3Qwen3VLConfig):
arch_config: QwenImage21Qwen3VLArchConfig = field(default_factory=QwenImage21Qwen3VLArchConfig)
prefix: str = "qwen_image21_qwen3_vl"
def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
flattened = dict(source_model_dict)
flattened.update(dict(flattened.pop("text_config", {})))
vision = flattened.pop("vision_config", {})
for source_name, target_name in _VISION_CONFIG_MAPPING.items():
if source_name in vision:
flattened[target_name] = vision[source_name]
architectures = flattened.get("architectures", [])
if isinstance(architectures, str):
architectures = [architectures]
if any(architecture in _OFFICIAL_ARCHITECTURES for architecture in architectures):
flattened["architectures"] = ["QwenImage21Qwen3VLConditioner"]
if "num_hidden_layers" in flattened:
flattened["output_hidden_state_index"] = flattened["num_hidden_layers"]
for name in ("mrope_section", "vision_deepstack_visual_indexes"):
if isinstance(flattened.get(name), list):
flattened[name] = tuple(flattened[name])
# Transformers 5 serializes the same default RoPE parameters under a
# renamed field. Keep the checkpoint's frequencies in either format.
if "rope_parameters" in flattened:
parameters = dict(flattened["rope_parameters"])
flattened["rope_scaling"] = parameters
if "rope_theta" in parameters:
flattened["rope_theta"] = parameters["rope_theta"]
ModelConfig.update_model_arch(self, flattened)
fastvideo/configs/models/vaes/__init__.py
+3
View File
@@ -12,6 +12,7 @@ from fastvideo.configs.models.vaes.minimax_h3_video import (
)
from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEArchConfig, QwenImage21VAEConfig
from fastvideo.models.wan.vae_config import WanVAEConfig
__all__ = [
@@ -31,4 +32,6 @@ __all__ = [
"OobleckVAEArchConfig",
"OobleckVAEConfig",
"Flux2VAEConfig",
"QwenImage21VAEArchConfig",
"QwenImage21VAEConfig",
]
fastvideo/configs/models/vaes/qwen_image21.py
+187
View File
@@ -0,0 +1,187 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 RGBA autoencoder configuration."""
from dataclasses import dataclass, field
from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class QwenImage21VAEArchConfig(VAEArchConfig):
_class_name: str = "AutoencoderKLQwenImage21"
base_dim: int = 96
decoder_base_dim: int | None = 144
z_dim: int = 64
dim_mult: tuple[int, ...] = (1, 2, 4, 8, 8)
num_res_blocks: int = 2
attn_scales: tuple[float, ...] = ()
temperal_downsample: tuple[bool, ...] = (False, True, True, True)
dropout: float = 0.0
is_residual: bool = True
in_channels: int = 4
out_channels: int = 4
patch_size: int | None = None
scale_factor_spatial: int = 16
scale_factor_temporal: int = 8
latents_mean: tuple[float, ...] = (
0.5126,
0.7721,
-0.0631,
1.3506,
-0.7855,
-2.1025,
-0.3458,
1.3722,
1.8873,
-1.7177,
-0.6510,
0.2732,
0.7562,
-0.6163,
-1.0277,
3.8363,
2.0210,
0.0472,
0.9320,
2.0087,
2.4954,
-0.1391,
-1.4249,
1.8464,
-0.5236,
1.2826,
3.7046,
-1.3035,
2.7286,
-1.4518,
-1.9036,
-1.9955,
-0.0342,
-1.0265,
-0.7636,
3.0555,
0.0746,
-3.0751,
-0.1076,
1.7376,
-1.0914,
-1.9435,
-0.2784,
-1.3680,
0.4809,
-0.4433,
0.3764,
0.5729,
-2.0595,
1.0960,
-1.3260,
-2.0211,
-5.0179,
0.5275,
4.0162,
1.8505,
0.3026,
1.9373,
1.4937,
0.2632,
0.5547,
-1.7121,
-0.1562,
0.0304,
)
latents_std: tuple[float, ...] = (
3.2001,
3.2936,
3.4321,
3.0091,
3.1061,
4.0379,
4.0705,
3.7910,
3.0785,
3.6500,
3.9308,
3.0904,
2.8778,
3.7675,
3.7320,
5.0756,
3.2864,
4.0397,
3.1317,
4.0443,
2.9249,
3.9454,
3.0988,
4.2489,
3.4896,
3.8513,
3.9323,
3.4719,
3.7498,
4.2830,
3.5694,
4.2467,
3.9037,
3.2947,
5.0770,
3.5075,
3.2700,
3.4767,
2.8063,
5.1125,
3.5327,
4.7833,
3.1286,
4.1819,
3.8527,
3.8312,
3.5605,
4.3875,
3.9624,
4.0168,
3.5643,
4.0550,
5.5614,
4.2963,
4.4080,
3.4959,
3.8747,
3.7608,
3.5735,
3.1490,
3.7662,
3.6746,
3.4563,
3.8161,
)
scaling_factor: float = 1.0
temporal_compression_ratio: int = 1
spatial_compression_ratio: int = 16
param_names_mapping: dict[str, str] = field(default_factory=lambda: {r"^(.*)$": r"\1"})
@property
def latent_channels(self) -> int:
return self.z_dim
def __post_init__(self) -> None:
if len(self.dim_mult) < 2 or len(self.temperal_downsample) != len(self.dim_mult) - 1:
raise ValueError("`temperal_downsample` must have one entry per downsampling stage.")
if len(self.latents_mean) != self.z_dim or len(self.latents_std) != self.z_dim:
raise ValueError("Qwen-Image-2.1 latent mean/std must have one entry per latent channel.")
if any(value <= 0 for value in self.latents_std):
raise ValueError("Qwen-Image-2.1 latent standard deviations must be positive.")
if self.patch_size is not None and self.patch_size < 1:
raise ValueError("`patch_size` must be positive or None.")
actual_spatial_scale = 2**(len(self.dim_mult) - 1) * (self.patch_size or 1)
if self.scale_factor_spatial != actual_spatial_scale:
raise ValueError("`scale_factor_spatial` must match the VAE downsampling stages and patch size.")
self.spatial_compression_ratio = self.scale_factor_spatial
@dataclass
class QwenImage21VAEConfig(VAEConfig):
arch_config: QwenImage21VAEArchConfig = field(default_factory=QwenImage21VAEArchConfig)
use_tiling: bool = False
use_temporal_tiling: bool = False
use_parallel_tiling: bool = False
fastvideo/configs/pipelines/__init__.py
+2 -1
View File
@@ -14,6 +14,7 @@ from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
from fastvideo.registry import get_pipeline_config_cls_from_name
from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
@@ -26,5 +27,5 @@ __all__ = [
"Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
"Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
"LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
"get_pipeline_config_cls_from_name"
"QwenImage21PipelineConfig", "get_pipeline_config_cls_from_name"
]
fastvideo/configs/pipelines/base.py
+1
View File
@@ -64,6 +64,7 @@ class PipelineConfig:
vae_decode_precision: str | None = None
vae_tiling: bool = True
vae_sp: bool = True
output_channels: int = 3
# Image encoder configuration
image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
fastvideo/configs/pipelines/qwen_image21.py
+30
View File
@@ -0,0 +1,30 @@
# SPDX-License-Identifier: Apache-2.0
"""Native Qwen-Image-2.1 inference configuration."""
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
@dataclass
class QwenImage21PipelineConfig(PipelineConfig):
dit_config: QwenImage21Config = field(default_factory=QwenImage21Config)
vae_config: QwenImage21VAEConfig = field(default_factory=QwenImage21VAEConfig)
text_encoder_configs: tuple[QwenImage21Qwen3VLConfig,
...] = field(default_factory=lambda: (QwenImage21Qwen3VLConfig(), ))
text_encoder_precisions: tuple[str, ...] = ("bf16", )
dit_precision: str = "bf16"
vae_precision: str = "bf16"
vae_tiling: bool = True
vae_sp: bool = False
# Four output channels are required to preserve transparency through PNG export.
output_channels: int = 4
flow_shift: float | None = None
scheduler_step_in_fp32: bool = True
reference_resolution: int = 1024
use_kv_cache: bool = True
# Runtime placement policy. CPU keeps ten reference prefixes out of VRAM.
kv_cache_device: str = "cpu"
fastvideo/entrypoints/video_generator.py
+2 -1
View File
@@ -949,7 +949,8 @@ class VideoGenerator:
samples = torch.empty(0, device='cpu')
else:
samples = allocate_cpu_tensor_with_pin_fallback(
(latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
(latent_batch_size, getattr(fastvideo_args.pipeline_config, "output_channels",
3), sampling_param.num_frames, sampling_param.height, sampling_param.width),
pin_memory=fastvideo_args.pin_cpu_memory)
try:
thread.join()
fastvideo/models/dits/qwen_image21.py
+592
View File
@@ -0,0 +1,592 @@
# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Native Qwen-Image-2.1 single-stream diffusion transformer.
The interleaved text/image block-causal stream has no matching FastVideo
sequence-parallel attention primitive. Mask-free dense calls use FastVideo's
shared attention selector, retaining PyTorch SDPA for causal text and padding
masks. Tensor and sequence parallelism are not implemented.
Prefix K/V can reside on the host, so decoding restores only one layer's cache
at a time instead of keeping every layer's reference images on the GPU.
"""
from __future__ import annotations
import math
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.attention.selector import effective_attention_backend
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21Config
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.logger import init_logger
from fastvideo.models.dits.base import BaseDiT
if TYPE_CHECKING:
from fastvideo.platforms import AttentionBackendEnum
logger = init_logger(__name__)
_IMG_TOKENS_PER_SLOT = 4
class QwenImage21KVCacheAllocationError(RuntimeError):
"""CPU prefix-cache allocation failed; an inference caller may retry uncached."""
class QwenImage21KVLayerCache:
"""Post-RoPE prefix K/V in ``(batch, tokens, heads, head_dim)`` format."""
def __init__(self, storage_device: str | torch.device | None = None):
self.storage_device = torch.device(storage_device) if storage_device is not None else None
self.k: torch.Tensor | None = None
self.v: torch.Tensor | None = None
def store(self, k: torch.Tensor, v: torch.Tensor) -> None:
if k.shape != v.shape or k.ndim != 4:
raise ValueError("Prefix K/V must have matching (batch, tokens, heads, head_dim) shapes")
# A fresh allocation releases the larger prefill tensor after the layer finishes.
device = self.storage_device or k.device
try:
self.k = k.detach().to(device=device, copy=True)
self.v = v.detach().to(device=device, copy=True)
except (MemoryError, RuntimeError) as exc:
self.clear()
message = str(exc).lower()
allocation_error = isinstance(exc, MemoryError) or any(
marker in message for marker in (
"defaultcpuallocator", "cannot allocate memory", "can't allocate memory",
"not enough memory", "out of memory", "std::bad_alloc"
))
if device.type == "cpu" and allocation_error and "cuda" not in message and "hip" not in message:
raise QwenImage21KVCacheAllocationError(
"CPU memory exhausted while allocating the prefix KV cache") from exc
raise
def get(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
if self.k is None or self.v is None:
raise RuntimeError("Prefix KV cache has not been populated; run an extract step first")
return self.k.to(device=device, dtype=dtype), self.v.to(device=device, dtype=dtype)
def clear(self) -> None:
self.k = self.v = None
class QwenImage21KVCache:
"""Per-request cache container; never share it between prompts or CFG branches."""
def __init__(self, num_layers: int, storage_device: str | torch.device | None = None):
if num_layers <= 0:
raise ValueError("num_layers must be positive")
self.layer_caches = [QwenImage21KVLayerCache(storage_device) for _ in range(num_layers)]
self.prefix_len: int | None = None
def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache:
return self.layer_caches[layer_idx]
def clear(self) -> None:
for layer in self.layer_caches:
layer.clear()
self.prefix_len = None
def apply_rotary_emb_qwen(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
return torch.view_as_real(rotated * freqs_cis.unsqueeze(1)).flatten(3).to(x.dtype)
class QwenImage21TemporalTimesteps(nn.Module):
def __init__(self, timestep_dim: int = 256, max_period: int = 10000, time_factor: float = 1000.0):
super().__init__()
self.timestep_dim = timestep_dim
self.time_factor = time_factor
half = timestep_dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(half, dtype=torch.float32, device="cpu") / half)
self.register_buffer("freqs", freqs, persistent=False)
def forward(self, timestep: torch.Tensor) -> torch.Tensor:
args = (self.time_factor * timestep.float())[:, None] * self.freqs.float()[None].to(timestep.device)
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if self.timestep_dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
return embedding.to(timestep.dtype)
class QwenImage21TimestepEmbedding(nn.Module):
def __init__(self, embedding_dim: int):
super().__init__()
self.linear_1 = ReplicatedLinear(256, embedding_dim, bias=False)
self.act = nn.SiLU()
self.linear_2 = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.linear_1(hidden_states)
return self.linear_2(self.act(hidden_states))[0]
class QwenImage21TimestepProjEmbeddings(nn.Module):
def __init__(self, embedding_dim: int):
super().__init__()
self.time_proj = QwenImage21TemporalTimesteps()
self.timestep_embedder = QwenImage21TimestepEmbedding(embedding_dim)
def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
return self.timestep_embedder(self.time_proj(timestep).to(hidden_states.dtype))
class QwenImage21ZeroCenterRMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.zeros(dim))
self.eps = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.float()
rrms = torch.rsqrt(torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps)
return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype)
class QwenImage21RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
variance = hidden_states.float().pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
if self.weight.dtype in (torch.float16, torch.bfloat16):
hidden_states = hidden_states.to(self.weight.dtype)
return hidden_states * self.weight
class QwenImage21TextProjection(nn.Module):
def __init__(self, context_in_dim: int, hidden_size: int, eps: float):
super().__init__()
self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps)
self.in_layer = ReplicatedLinear(context_in_dim, hidden_size, bias=False)
self.act = nn.GELU(approximate="tanh")
self.out_layer = ReplicatedLinear(hidden_size, hidden_size, bias=False)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.in_layer(self.text_norm(hidden_states))
return self.out_layer(self.act(hidden_states))[0]
class QwenImage21SwiGLUFeedForward(nn.Module):
def __init__(self, hidden_size: int, mlp_hidden_size: int):
super().__init__()
self.proj = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
self.out = ReplicatedLinear(mlp_hidden_size, hidden_size, bias=False)
self.gate_layer = ReplicatedLinear(hidden_size, mlp_hidden_size, bias=False)
self.activation_fn = nn.SiLU()
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
gate = self.activation_fn(self.gate_layer(hidden_states)[0])
return self.out(gate * self.proj(hidden_states)[0])[0]
def _select_modulation_rows(params: torch.Tensor, target_token_mask: torch.Tensor | None) -> torch.Tensor:
if target_token_mask is None:
return params.unsqueeze(1)
real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0)
return torch.where(target_token_mask.view(1, -1, 1), real, zero)
class QwenImage21AdaLayerNormContinuous(nn.Module):
def __init__(self, embedding_dim: int, eps: float):
super().__init__()
self.silu = nn.SiLU()
self.linear = ReplicatedLinear(embedding_dim, embedding_dim, bias=False)
self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=False)
def forward(self, hidden_states, conditioning_embedding, target_token_mask=None):
scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype))[0]
return self.norm(hidden_states) * (1 + _select_modulation_rows(scale, target_token_mask))
def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]:
prefix_ids = image_ids[:prefix_len].tolist()
segments = []
start = 0
for index in range(1, prefix_len + 1):
if index == prefix_len or prefix_ids[index] != prefix_ids[start]:
segments.append((start, index, prefix_ids[start] < 0))
start = index
return segments
def _attention(query, key, value, mask=None):
# SDPA expects (batch, heads, tokens, dim); projected QKV use tokens before heads.
output = F.scaled_dot_product_attention(query.transpose(1, 2),
key.transpose(1, 2),
value.transpose(1, 2),
attn_mask=mask,
dropout_p=0.0)
if mask is not None:
# Older CPU SDPA returns NaNs for fully masked padding queries.
output = output.masked_fill(~mask.any(dim=-1, keepdim=True), 0.0)
return output.transpose(1, 2)
def _make_dense_attention(heads: int, dim_head: int,
supported_attention_backends: tuple[AttentionBackendEnum, ...],
requested_backend: AttentionBackendEnum | None) -> nn.Module | None:
# Keep optional attention kernels and GPU dispatch imports off the SDPA path.
from fastvideo.attention.layer import LocalAttention
from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
from fastvideo.platforms import AttentionBackendEnum
from fastvideo.utils import get_compute_dtype
dtype = get_compute_dtype()
selection_kwargs = {} if requested_backend is None else {"requested": requested_backend}
backend_cls = get_attn_backend(dim_head, dtype,
supported_attention_backends=supported_attention_backends,
**selection_kwargs)
backend = backend_name_to_enum(backend_cls.get_name())
if backend is AttentionBackendEnum.TORCH_SDPA:
return None
class DenseLocalAttention(LocalAttention):
"""Qwen owns segment lengths and has already applied its complex RoPE.
Resolve explicit component requests through the public selector above;
direct construction without a request retains scope/env selection.
LocalAttention's ordinary forward reads unrelated global metadata,
which must not reach our unequal Q/KV calls.
"""
def __init__(self):
nn.Module.__init__(self)
self.softmax_scale = dim_head**-0.5
self.attn_impl = backend_cls.get_impl_cls()(num_heads=heads, head_size=dim_head,
softmax_scale=self.softmax_scale,
num_kv_heads=heads, causal=False)
self.num_heads = self.num_kv_heads = heads
self.head_size, self.backend, self.dtype = dim_head, backend, dtype
def forward(self, q, k, v):
return self.attn_impl.forward(q, k, v, None)
return DenseLocalAttention()
class QwenImage21Attention(nn.Module):
def __init__(self, dim: int, heads: int, dim_head: int, eps: float, *,
requested_backend: AttentionBackendEnum | None = None,
supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
super().__init__()
self.heads = heads
self.inner_dim = heads * dim_head
self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=False)
self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=False)
self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=False)
self.to_out = nn.ModuleList([ReplicatedLinear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)])
self.norm_q = QwenImage21RMSNorm(dim_head, eps)
self.norm_k = QwenImage21RMSNorm(dim_head, eps)
self.dense_attn = None
if requested_backend is None or requested_backend.name != "TORCH_SDPA":
supported = supported_attention_backends or _DEFAULT_CONFIG.arch_config._supported_attention_backends
self.dense_attn = _make_dense_attention(heads, dim_head, supported, requested_backend)
def _attend(self, query, key, value, mask=None):
if mask is not None or self.dense_attn is None:
return _attention(query, key, value, mask)
return self.dense_attn(query, key, value)
def forward(self,
hidden_states: torch.Tensor,
rotary_emb: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
layer_cache: QwenImage21KVLayerCache | None = None,
kv_cache_mode: str | None = None,
cache_write_slice: slice | None = None,
segments: list[tuple[int, int, bool]] | None = None,
key_valid: torch.Tensor | None = None) -> torch.Tensor:
query = self.to_q(hidden_states)[0].unflatten(-1, (self.heads, -1))
key = self.to_k(hidden_states)[0].unflatten(-1, (self.heads, -1))
value = self.to_v(hidden_states)[0].unflatten(-1, (self.heads, -1))
query, key = self.norm_q(query).to(value.dtype), self.norm_k(key).to(value.dtype)
if rotary_emb is not None:
query = apply_rotary_emb_qwen(query, rotary_emb)
key = apply_rotary_emb_qwen(key, rotary_emb)
if layer_cache is not None:
if kv_cache_mode == "extract" and cache_write_slice is not None:
layer_cache.store(key[:, cache_write_slice], value[:, cache_write_slice])
elif kv_cache_mode == "cached":
cached_k, cached_v = layer_cache.get(key.device, key.dtype)
if cached_k.shape[0] != key.shape[0] or cached_k.shape[2:] != key.shape[2:]:
raise ValueError("Cached prefix batch/head shape differs from the current request")
key = torch.cat([cached_k, key], dim=1)
value = torch.cat([cached_v, value], dim=1)
del cached_k, cached_v
if segments is None:
output = self._attend(query, key, value, attention_mask)
else:
outputs = []
for start, end, is_text in segments:
mask = None
if is_text:
length = end - start
mask = torch.cat([torch.ones(length, start, dtype=torch.bool, device=query.device),
torch.tril(torch.ones(length, length, dtype=torch.bool, device=query.device))],
dim=1)[None, None]
if key_valid is not None:
valid = key_valid[:, None, None, :end]
mask = valid if mask is None else mask & valid
outputs.append(self._attend(query[:, start:end], key[:, :end], value[:, :end], mask))
prefix_len = segments[-1][1] if segments else 0
mask = None if key_valid is None else key_valid[:, None, None, :]
outputs.append(self._attend(query[:, prefix_len:], key, value, mask))
output = torch.cat(outputs, dim=1)
output = output.flatten(2, 3).to(query.dtype)
return self.to_out[1](self.to_out[0](output)[0])
class QwenImage21TransformerBlock(nn.Module):
def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: int, eps: float, *,
requested_backend: AttentionBackendEnum | None = None,
supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None):
super().__init__()
self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
self.attn = QwenImage21Attention(dim, num_attention_heads, attention_head_dim, eps,
requested_backend=requested_backend,
supported_attention_backends=supported_attention_backends)
self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
self.img_mlp = QwenImage21SwiGLUFeedForward(dim, dim * mlp_ratio)
def _modulate(self, hidden_states, mod_params, target_token_mask):
scale, gate = mod_params.chunk(2, dim=-1)
scale = _select_modulation_rows(scale, target_token_mask)
gate = _select_modulation_rows(gate, target_token_mask)
return hidden_states * (1 + scale), gate
def forward(self, hidden_states, modulation, target_token_mask=None, **attention_kwargs):
mod1, mod2 = modulation.chunk(2, dim=-1)
normalized, gate = self._modulate(self.img_norm1(hidden_states), mod1, target_token_mask)
hidden_states = hidden_states + gate.tanh() * self.attn(normalized, **attention_kwargs)
normalized, gate = self._modulate(self.img_norm2(hidden_states), mod2, target_token_mask)
hidden_states = hidden_states + gate.tanh() * self.img_mlp(normalized)
if hidden_states.dtype == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
return hidden_states
class QwenImage21Rope(nn.Module):
def __init__(self, theta: int, axes_dim: tuple[int, int, int]):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim
def rope_params(self, index: torch.Tensor, dim: int) -> torch.Tensor:
powers = torch.arange(0, dim, 2, dtype=torch.float32, device=index.device).div(dim)
freqs = torch.outer(index.float(), 1.0 / torch.pow(self.theta, powers))
return torch.polar(torch.ones_like(freqs), freqs)
def forward(self, img_shapes, image_pad_mask, device):
frame_index, image_height_index, image_width_index = [], [], []
cursor, position = 0, 0
is_image = image_pad_mask.tolist()
for _, height, width in img_shapes:
block_start = is_image.index(True, cursor)
text_len = block_start - cursor
frame_index.extend(range(position, position + text_len))
position += text_len
cursor = block_start + height * width
frame_index.extend([position] * (height * width))
position += max(height, width)
image_height_index.extend([h for h in range(-(height - height // 2), height // 2) for _ in range(width)])
image_width_index.extend([w for _ in range(height) for w in range(-(width - width // 2), width // 2)])
if cursor < len(is_image):
frame_index.extend(range(position, position + len(is_image) - cursor))
# Derive frequencies from signed coordinates, retaining Qwen's CPU fp32/complex64 numerics.
frame_index = torch.tensor(frame_index, dtype=torch.long, device="cpu")
image_pad_mask = torch.tensor(is_image, dtype=torch.bool, device="cpu")
height_index, width_index = frame_index.clone(), frame_index.clone()
height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device="cpu")
width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device="cpu")
return torch.cat([self.rope_params(index, dim) for index, dim in
zip((frame_index, height_index, width_index), self.axes_dim, strict=True)], dim=-1).to(device)
_DEFAULT_CONFIG = QwenImage21Config()
class QwenImage21Transformer2DModel(BaseDiT):
_fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
_compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
_supported_attention_backends = _DEFAULT_CONFIG.arch_config._supported_attention_backends
param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
reverse_param_names_mapping = _DEFAULT_CONFIG.arch_config.reverse_param_names_mapping
def __init__(self, config: QwenImage21Config, hf_config: dict[str, Any] | None = None, **kwargs):
super().__init__(config, hf_config or {}, **kwargs)
if getattr(config, "quant_config", None) is not None:
raise NotImplementedError("Qwen-Image-2.1 transformer quantization has not been implemented")
self.inner_dim = self.hidden_size = config.hidden_size
self.out_channels = config.out_channels
self.num_attention_heads = config.num_attention_heads
self.num_channels_latents = config.num_channels_latents
self.pos_embed = QwenImage21Rope(10000, config.axes_dims_rope)
self.time_text_embed = QwenImage21TimestepProjEmbeddings(self.inner_dim)
self.txt_in = QwenImage21TextProjection(config.context_in_dim, self.inner_dim, config.eps)
self.img_in = ReplicatedLinear(config.in_channels * config.patch_size**2, self.inner_dim, bias=False)
self.modulation = nn.Sequential(nn.SiLU(), ReplicatedLinear(self.inner_dim, 4 * self.inner_dim, bias=False))
requested_backend = effective_attention_backend(config)
self.transformer_blocks = nn.ModuleList([
QwenImage21TransformerBlock(self.inner_dim, config.num_attention_heads, config.attention_head_dim,
config.mlp_ratio, config.eps, requested_backend=requested_backend,
supported_attention_backends=self._supported_attention_backends)
for _ in range(config.num_layers)
])
self.norm_out = QwenImage21AdaLayerNormContinuous(self.inner_dim, config.eps)
self.proj_out = ReplicatedLinear(self.inner_dim, config.patch_size**2 * config.out_channels, bias=False)
@staticmethod
def build_token_metadata(image_pad_mask, img_shapes):
positions = image_pad_mask.nonzero(as_tuple=True)[0]
lengths = [math.prod(shape) for shape in img_shapes]
if not lengths or any(length <= 0 for length in lengths) or sum(lengths) != positions.numel():
raise ValueError("img_shapes token counts must match image_pad_mask")
image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long)
ids = torch.repeat_interleave(torch.arange(len(lengths), device=image_pad_mask.device),
torch.tensor(lengths, device=image_pad_mask.device))
image_ids[positions] = ids
target_mask = torch.zeros_like(image_pad_mask)
target_mask[positions[-lengths[-1]:]] = True
return image_ids, target_mask
def _validate_inputs(self, hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask, text_mask):
batch_size = hidden_states.shape[0]
if hidden_states.ndim != 3 or encoder_hidden_states.ndim != 3:
raise ValueError("Latents and text embeddings must have (batch, tokens, channels) shapes")
if encoder_hidden_states.shape[0] != batch_size:
raise ValueError("Latent and text batch sizes must match")
if len(img_shapes) != batch_size or not img_shapes or not img_shapes[0]:
raise ValueError("img_shapes must provide a nonempty image layout for each sample")
if any(shapes != img_shapes[0] for shapes in img_shapes):
raise ValueError("All samples must share the same image layout")
if any(frame != 1 or height <= 0 or width <= 0 for frame, height, width in img_shapes[0]):
raise ValueError("Qwen-Image-2.1 expects positive single-frame image shapes")
target_tokens = math.prod(img_shapes[0][-1])
if target_tokens % _IMG_TOKENS_PER_SLOT:
raise ValueError("The target latent token count must be divisible by four")
if hidden_states.shape[1] != sum(math.prod(shape) for shape in img_shapes[0]):
raise ValueError("Packed latents must include all references followed by the target image")
expected_slots = encoder_hidden_states.shape[1] + target_tokens // _IMG_TOKENS_PER_SLOT
if img_mask.shape != (batch_size, expected_slots):
raise ValueError("img_mask must include the VLM sequence and appended target-image slots")
if not torch.equal(img_mask.bool(), img_mask[0:1].bool().expand_as(img_mask)):
raise ValueError("All samples must share the same image-slot layout")
if not bool(img_mask[:, -target_tokens // _IMG_TOKENS_PER_SLOT:].all()):
raise ValueError("Target-image slots must be appended at the end of img_mask")
if int(img_mask[0].count_nonzero()) * _IMG_TOKENS_PER_SLOT != hidden_states.shape[1]:
raise ValueError("Image-slot count must match the reference and target latent token count")
if timestep.numel() != batch_size:
raise ValueError("timestep must contain one value per sample")
if text_mask is not None and text_mask.shape != encoder_hidden_states.shape[:2]:
raise ValueError("encoder_hidden_states_mask must match the VLM embedding sequence")
def forward(self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
timestep: torch.Tensor,
img_shapes: list[list[tuple[int, int, int]]],
img_mask: torch.Tensor,
encoder_hidden_states_mask: torch.Tensor | None = None,
attention_kwargs: dict[str, Any] | None = None,
kv_cache: QwenImage21KVCache | None = None,
kv_cache_mode: str | None = None,
return_dict: bool | None = None,
**kwargs) -> torch.Tensor | tuple[torch.Tensor]:
"""Return joint-sequence predictions; the pipeline selects the target tail.
``kv_cache_mode='extract'`` stores timestep-independent text/reference
K/V. ``'cached'`` returns only the target image tokens. By default the
return value is a tensor; ``return_dict=False`` selects a one-item tuple.
"""
self._validate_inputs(hidden_states, encoder_hidden_states, timestep, img_shapes, img_mask,
encoder_hidden_states_mask)
if kv_cache is not None and not self.config.causal_condition:
raise ValueError("Prefix KV caching requires causal_condition=True")
if kv_cache is not None and kv_cache_mode not in ("extract", "cached"):
raise ValueError("kv_cache_mode must be extract or cached when a cache is provided")
if kv_cache is None and kv_cache_mode is not None:
raise ValueError("kv_cache_mode requires a cache container")
if kv_cache is not None and len(kv_cache.layer_caches) != len(self.transformer_blocks):
raise ValueError("KV cache layer count differs from the transformer")
if attention_kwargs:
raise ValueError("Custom attention processors and LoRA scaling are not implemented for Qwen-Image-2.1")
batch_size = hidden_states.shape[0]
hidden_states = self.img_in(hidden_states)[0]
encoder_hidden_states = self.txt_in(encoder_hidden_states)
img_mask = img_mask.bool().to(hidden_states.device)
repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0]
image_pad_mask = torch.repeat_interleave(img_mask[0], repeats)
target_tokens = math.prod(img_shapes[0][-1])
joint = torch.cat([encoder_hidden_states,
encoder_hidden_states.new_zeros(batch_size, target_tokens // 4, self.inner_dim)], dim=1)
joint = joint.repeat_interleave(repeats, dim=1)
joint[:, image_pad_mask] = hidden_states
rotary_emb = self.pos_embed(img_shapes[0], image_pad_mask, hidden_states.device)
image_ids, target_mask = self.build_token_metadata(image_pad_mask, img_shapes[0])
timestep = timestep.reshape(batch_size).to(device=hidden_states.device, dtype=hidden_states.dtype)
if self.config.causal_condition:
timestep = torch.cat([timestep, timestep.new_zeros(1)])
modulation_mask = target_mask
else:
modulation_mask = None
temb = self.time_text_embed(timestep, hidden_states)
modulation = self.modulation[1](self.modulation[0](temb))[0]
joint_key_valid = None
if encoder_hidden_states_mask is not None:
joint_key_valid = torch.ones(batch_size, image_pad_mask.numel(), dtype=torch.bool,
device=hidden_states.device)
text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0]
vlm_text_positions = ~img_mask[0, :encoder_hidden_states.shape[1]]
valid_text = encoder_hidden_states_mask.bool().to(hidden_states.device)[:, vlm_text_positions]
joint_key_valid[:, text_positions] = valid_text
prefix_len = image_pad_mask.numel() - target_tokens
if kv_cache_mode == "cached":
if kv_cache.prefix_len != prefix_len:
raise ValueError("Cached prefix layout differs from the current request or has not been extracted")
joint, rotary_emb = joint[:, prefix_len:], rotary_emb[prefix_len:]
modulation_mask = modulation_mask[prefix_len:]
attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :]
segments, block_key_valid, cache_slice = None, None, None
else:
segments = _qwenimage21_prefix_segments(image_ids, prefix_len)
attention_mask, block_key_valid = None, joint_key_valid
cache_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None
if kv_cache is not None:
kv_cache.prefix_len = prefix_len
for index, block in enumerate(self.transformer_blocks):
joint = block(joint, modulation,
target_token_mask=modulation_mask,
rotary_emb=rotary_emb,
attention_mask=attention_mask,
layer_cache=kv_cache.get_layer(index) if kv_cache is not None else None,
kv_cache_mode=kv_cache_mode,
cache_write_slice=cache_slice,
segments=segments,
key_valid=block_key_valid)
output = self.proj_out(self.norm_out(joint, temb, modulation_mask))[0]
return (output,) if return_dict is False else output
EntryClass = QwenImage21Transformer2DModel
fastvideo/models/encoders/qwen_image21.py
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
"""Native Qwen3-VL conditioner returning Qwen-Image-2.1's pre-norm features."""
from collections.abc import Iterable
import torch
from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
from fastvideo.models.loader.weight_utils import default_weight_loader
class QwenImage21Qwen3VLConditioner(MiniMaxH3Qwen3VLConditioner):
"""Qwen3-VL text and image features, without autoregressive generation.
The shared native stack stops immediately after the configured last
decoder layer. Its final RMSNorm is loaded for checkpoint completeness
but never applied to the image transformer's conditioning features.
Sequence parallel execution is not implemented for this conditioner.
"""
supported_checkpoint_quantization_methods: frozenset[str] = frozenset()
_checkpoint_quantization_configs: dict = {}
def __init__(self, config: QwenImage21Qwen3VLConfig) -> None:
super().__init__(config)
# CPU construction of this constant avoids device-dependent power
# rounding accumulating across the complete 36-layer decoder stack.
rotary = self.language_model.rotary_emb
exponents = torch.arange(0, config.head_dim, 2, dtype=torch.float32, device="cpu") / config.head_dim
rotary.inv_freq = (1.0 / (config.rope_theta**exponents)).to(rotary.inv_freq.device)
def _visual_features(
self,
pixels: torch.Tensor,
grid_thw: torch.Tensor,
) -> tuple[torch.Tensor, list[torch.Tensor]]:
if grid_thw.ndim != 2 or grid_thw.shape[1] != 3 or grid_thw.shape[0] == 0:
raise ValueError("image_grid_thw must contain one nonempty [time, height, width] row per image")
grids = grid_thw.to(device="cpu", dtype=torch.long)
merge = self.config.vision_spatial_merge_size
if bool((grids <= 0).any()) or bool((grids[:, 1:] % merge != 0).any()):
raise ValueError("Image grids must be positive and spatially divisible by the vision merge size")
patch_counts = grids.prod(dim=1).tolist()
if sum(patch_counts) != pixels.shape[0]:
raise ValueError("pixel_values patch count does not match image_grid_thw")
# Keep packed projection shapes: splitting images changes BF16 GEMM
# rounding before the complete text decoder amplifies the difference.
# Vision attention still executes separately for each image grid.
return self.visual(pixels.to(self.visual.patch_embed.proj.weight.dtype), grid_thw)
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
*,
attention_mask: torch.Tensor | None = None,
pixel_values: torch.Tensor | None = None,
image_grid_thw: torch.Tensor | None = None,
mm_token_type_ids: torch.Tensor | None = None,
) -> torch.Tensor:
"""Return ``[batch, tokens, width]``, or ``[tokens, width]`` for 1-D IDs.
Padding is zeroed and removed by the prompt stage using the original
attention mask before stripping its template.
"""
single_sequence = input_ids.ndim == 1
if single_sequence:
input_ids = input_ids.unsqueeze(0)
if attention_mask is not None and attention_mask.ndim == 1:
attention_mask = attention_mask.unsqueeze(0)
if mm_token_type_ids is not None and mm_token_type_ids.ndim == 1:
mm_token_type_ids = mm_token_type_ids.unsqueeze(0)
if input_ids.ndim != 2 or input_ids.shape[1] == 0:
raise ValueError("Qwen-Image-2.1 input_ids must have shape [batch, nonempty sequence] or [sequence]")
if attention_mask is not None:
if attention_mask.shape != input_ids.shape:
raise ValueError("attention_mask must have the same shape as input_ids")
if bool((attention_mask.to(torch.bool).sum(dim=1) == 0).any()):
raise ValueError("Each conditioning sequence must contain an unmasked token")
if (pixel_values is None) != (image_grid_thw is None):
raise ValueError("pixel_values and image_grid_thw must be provided together")
if bool((input_ids == self.config.video_token_id).any()):
raise ValueError("Qwen-Image-2.1 conditioning accepts image references, not video tokens")
if mm_token_type_ids is not None:
expected = (input_ids == self.config.image_token_id).to(mm_token_type_ids.dtype)
if mm_token_type_ids.shape != input_ids.shape or not torch.equal(mm_token_type_ids, expected):
raise ValueError("mm_token_type_ids must mark image placeholders as 1 and other tokens as 0")
inputs_embeds = self.language_model.embed_tokens(input_ids)
visual_mask = None
deepstack = None
if pixel_values is not None:
features, deepstack = self._visual_features(pixel_values, image_grid_thw)
features = features.to(inputs_embeds.device, inputs_embeds.dtype)
visual_mask = self._placeholder_mask(input_ids, inputs_embeds, self.config.image_token_id,
features, "image")
inputs_embeds = inputs_embeds.masked_scatter(visual_mask.unsqueeze(-1), features)
elif bool((input_ids == self.config.image_token_id).any()):
raise ValueError("Image placeholder tokens require pixel_values and image_grid_thw")
# Text-only prefills retain positions in the padded sequence. A
# constant RoPE shift is mathematically equivalent, but changes BF16
# rounding; image prefills need each sample's multimodal layout.
positions = self._get_rope_index(input_ids, image_grid_thw, None,
attention_mask if image_grid_thw is not None else None)
# Unpadding changes projection and attention shapes in reduced
# precision. Preserve the batch through every decoder layer instead.
hidden_states = self.language_model(inputs_embeds, positions, attention_mask, visual_mask, deepstack)
if attention_mask is not None:
hidden_states = hidden_states.masked_fill(~attention_mask.to(torch.bool).unsqueeze(-1), 0)
return hidden_states[0] if single_sequence else hidden_states
@torch.no_grad()
def encode_ids(
self,
input_ids: torch.Tensor,
*,
attention_mask: torch.Tensor | None = None,
pixel_values: torch.Tensor | None = None,
image_grid_thw: torch.Tensor | None = None,
mm_token_type_ids: torch.Tensor | None = None,
) -> torch.Tensor:
"""Encode one sequence without imposing a tokenizer truncation limit."""
if input_ids.ndim != 1:
raise ValueError("encode_ids requires 1-D input_ids; use forward for a padded batch")
return self.forward(
input_ids,
attention_mask=attention_mask,
pixel_values=pixel_values,
image_grid_thw=image_grid_thw,
mm_token_type_ids=mm_token_type_ids,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
parameters = dict(self.named_parameters())
loaded: set[str] = set()
for source_name, tensor in weights:
if source_name == "lm_head.weight":
continue
name = source_name.removeprefix("model.")
if name not in parameters:
raise ValueError(f"Unexpected Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
if name in loaded:
raise ValueError(f"Duplicate Qwen-Image-2.1 Qwen3-VL checkpoint key: {source_name}")
parameter = parameters[name]
loader = getattr(parameter, "weight_loader", default_weight_loader)
loader(parameter, tensor)
loaded.add(name)
return loaded
EntryClass = QwenImage21Qwen3VLConditioner
fastvideo/models/loader/component_loader.py
+1 -1
View File
@@ -1032,7 +1032,7 @@ class VAELoader(ComponentLoader):
# Diffusers-format AutoencoderKL checkpoints should match exactly; load
# strictly so missing/unexpected keys are surfaced early.
# Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE", "AutoencoderKLQwenImage21"}
vae.load_state_dict(loaded, strict=strict_load)
del loaded
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
fastvideo/models/registry.py
+7
View File
@@ -71,7 +71,12 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
"MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
}
_TEXT_TO_IMAGE_DIT_MODELS = {
"QwenImage21Transformer2DModel": ("dits", "qwen_image21", "QwenImage21Transformer2DModel"),
}
_TEXT_ENCODER_MODELS = {
"QwenImage21Qwen3VLConditioner": ("encoders", "qwen_image21", "QwenImage21Qwen3VLConditioner"),
"MMAudioDFNCLIPTextEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPTextEncoder"),
"CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
"CLIPTextModelWithProjection": ("encoders", "clip", "CLIPTextModelWithProjection"),
@@ -108,6 +113,7 @@ _VAE_MODELS = {
"AutoencoderKLWan": ("wan", "vae", "AutoencoderKLWan"),
"LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
"AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
"AutoencoderKLQwenImage21": ("vaes", "qwen_image21", "AutoencoderKLQwenImage21"),
"AutoencoderKLGen3CTokenizer": ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
"CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
"AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
@@ -154,6 +160,7 @@ _UPSAMPLERS = {
_LEGACY_FAST_VIDEO_MODELS = {
**_TEXT_TO_VIDEO_DIT_MODELS,
**_IMAGE_TO_VIDEO_DIT_MODELS,
**_TEXT_TO_IMAGE_DIT_MODELS,
**_TEXT_ENCODER_MODELS,
**_IMAGE_ENCODER_MODELS,
**_VAE_MODELS,
fastvideo/models/vaes/qwen_image21.py
+656
View File
@@ -0,0 +1,656 @@
# Copyright 2026 The Qwen 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.
"""Native single-frame RGBA autoencoder for Qwen-Image-2.1."""
from __future__ import annotations
import json
import math
from dataclasses import dataclass
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F
from fastvideo.configs.models.vaes.base import VAEConfig
from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEConfig
from fastvideo.logger import init_logger
logger = init_logger(__name__)
class DiagonalGaussianDistribution:
def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None:
self.parameters = parameters
self.mean, self.logvar = parameters.chunk(2, dim=1)
self.logvar = self.logvar.clamp(-30.0, 20.0)
self.deterministic = deterministic
self.std = torch.exp(0.5 * self.logvar)
self.var = torch.exp(self.logvar)
if deterministic:
self.std = torch.zeros_like(self.mean)
self.var = torch.zeros_like(self.mean)
def mode(self) -> torch.Tensor:
return self.mean
def sample(self, generator: torch.Generator | list[torch.Generator] | None = None) -> torch.Tensor:
if isinstance(generator, list):
if len(generator) != self.mean.shape[0]:
raise ValueError("A generator list must contain one generator per batch element.")
noise = torch.cat([
self._noise(self.mean[index:index + 1], item) for index, item in enumerate(generator)
])
else:
noise = self._noise(self.mean, generator)
return self.mean + self.std * noise
@staticmethod
def _noise(sample: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
noise_device = sample.device if generator is None else generator.device
return torch.randn(sample.shape, generator=generator, device=noise_device, dtype=sample.dtype).to(sample.device)
def kl(self, other: DiagonalGaussianDistribution | None = None) -> torch.Tensor:
if self.deterministic:
return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
if other is None:
terms = self.mean.square() + self.var - 1.0 - self.logvar
else:
terms = ((self.mean - other.mean).square() + self.var) / other.var - 1.0 - self.logvar + other.logvar
return 0.5 * terms.sum(dim=(1, 2, 3))
def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor:
if self.deterministic:
return torch.zeros(1, device=self.mean.device, dtype=self.mean.dtype)
return 0.5 * (math.log(2 * math.pi) + self.logvar + (sample - self.mean).square() / self.var).sum(dim=dims)
@dataclass
class AutoencoderKLOutput:
latent_dist: DiagonalGaussianDistribution
@dataclass
class DecoderOutput:
sample: torch.Tensor
class QwenImage21AvgDown3D(nn.Module):
def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.factor_t = factor_t
self.factor_s = factor_s
self.factor = factor_t * factor_s * factor_s
if in_channels * self.factor % out_channels:
raise ValueError("Average shortcut channels must be divisible by the downsampling factor.")
self.group_size = in_channels * self.factor // out_channels
def forward(self, x: torch.Tensor) -> torch.Tensor:
pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
batch, channels, frames, height, width = x.shape
x = x.view(batch, channels, frames // self.factor_t, self.factor_t,
height // self.factor_s, self.factor_s, width // self.factor_s, self.factor_s)
x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
x = x.view(batch, channels * self.factor, frames // self.factor_t,
height // self.factor_s, width // self.factor_s)
x = x.view(batch, self.out_channels, self.group_size, frames // self.factor_t,
height // self.factor_s, width // self.factor_s)
return x.mean(dim=2)
class QwenImage21DupUp3D(nn.Module):
def __init__(self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1) -> None:
super().__init__()
self.out_channels = out_channels
self.factor_t = factor_t
self.factor_s = factor_s
factor = factor_t * factor_s * factor_s
if out_channels * factor % in_channels:
raise ValueError("Duplicate shortcut channels must be divisible by the upsampling factor.")
self.repeats = out_channels * factor // in_channels
def forward(self, x: torch.Tensor, first_chunk: bool = False) -> torch.Tensor:
x = x.repeat_interleave(self.repeats, dim=1)
x = x.view(x.shape[0], self.out_channels, self.factor_t, self.factor_s, self.factor_s,
x.shape[2], x.shape[3], x.shape[4])
x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
x = x.view(x.shape[0], self.out_channels, x.shape[2] * self.factor_t,
x.shape[4] * self.factor_s, x.shape[6] * self.factor_s)
return x[:, :, self.factor_t - 1:] if first_chunk else x
class QwenImage21CausalConv3d(nn.Conv2d):
"""Spatial convolution retaining the image VAE's singleton frame axis."""
def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0) -> None:
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
self._padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0])
self.padding = (0, 0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.ndim != 5 or x.shape[2] != 1:
raise ValueError("Qwen-Image-2.1 spatial convolutions require exactly one frame.")
return super().forward(F.pad(x.squeeze(2), self._padding)).unsqueeze(2)
class QwenImage21RMSNorm(nn.Module):
def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
super().__init__()
broadcast_dims = (1, 1) if images else (1, 1, 1)
shape = (dim, *broadcast_dims) if channel_first else (dim,)
self.channel_first = channel_first
self.scale = dim**0.5
self.gamma = nn.Parameter(torch.ones(shape))
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
def forward(self, x: torch.Tensor) -> torch.Tensor:
upcast = x.dtype in (torch.float16, torch.bfloat16) or any(
name in str(x.dtype) for name in ("float4_", "float8_")
)
normalized = F.normalize(x.float() if upcast else x, dim=1 if self.channel_first else -1).to(x.dtype)
return normalized * self.scale * self.gamma + self.bias
class QwenImage21Upsample(nn.Upsample):
def forward(self, x: torch.Tensor) -> torch.Tensor:
return super().forward(x.float()).type_as(x)
class QwenImage21Resample(nn.Module):
def __init__(self, dim: int, mode: str, upsample_out_dim: int | None = None) -> None:
super().__init__()
self.mode = mode
if upsample_out_dim is None:
upsample_out_dim = dim // 2
if mode in ("upsample2d", "upsample3d"):
self.resample = nn.Sequential(
QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
)
if mode == "upsample3d":
self.time_conv = QwenImage21CausalConv3d(dim, dim * 2, (1, 1), padding=(0, 0))
elif mode in ("downsample2d", "downsample3d"):
self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
if mode == "downsample3d":
self.time_conv = QwenImage21CausalConv3d(dim, dim, (1, 1), stride=(1, 1), padding=(0, 0))
else:
self.resample = nn.Identity()
def forward(self, x: torch.Tensor) -> torch.Tensor:
batch, channels, frames, height, width = x.shape
# Temporal convolutions have no preceding frame in this image model.
x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
x = self.resample(x)
return x.view(batch, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
class QwenImage21ResidualBlock(nn.Module):
def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
super().__init__()
self.norm1 = QwenImage21RMSNorm(in_dim, images=False)
self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1)
self.norm2 = QwenImage21RMSNorm(out_dim, images=False)
self.dropout = nn.Dropout(dropout)
self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1)
self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = self.conv_shortcut(x)
x = self.conv1(F.silu(self.norm1(x)))
x = self.conv2(self.dropout(F.silu(self.norm2(x))))
return x + residual
class QwenImage21AttentionBlock(nn.Module):
def __init__(self, dim: int) -> None:
super().__init__()
self.norm = QwenImage21RMSNorm(dim)
self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
self.proj = nn.Conv2d(dim, dim, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
batch, channels, frames, height, width = x.shape
x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
qkv = self.to_qkv(self.norm(x)).reshape(batch * frames, 1, channels * 3, -1)
q, k, v = qkv.permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
# The VAE's one-head spatial attention has no distributed sequence layout.
x = F.scaled_dot_product_attention(q, k, v)
x = x.squeeze(1).permute(0, 2, 1).reshape(batch * frames, channels, height, width)
x = self.proj(x).view(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4)
return x + residual
class QwenImage21MidBlock(nn.Module):
def __init__(self, dim: int, dropout: float = 0.0) -> None:
super().__init__()
self.resnets = nn.ModuleList([
QwenImage21ResidualBlock(dim, dim, dropout),
QwenImage21ResidualBlock(dim, dim, dropout),
])
self.attentions = nn.ModuleList([QwenImage21AttentionBlock(dim)])
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.resnets[1](self.attentions[0](self.resnets[0](x)))
class QwenImage21ResidualDownBlock(nn.Module):
def __init__(self, in_dim: int, out_dim: int, dropout: float, num_res_blocks: int,
temporal_downsample: bool, down_flag: bool) -> None:
super().__init__()
self.avg_shortcut = QwenImage21AvgDown3D(in_dim, out_dim, 2 if temporal_downsample else 1,
2 if down_flag else 1)
self.resnets = nn.ModuleList([
QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, dropout)
for index in range(num_res_blocks)
])
self.downsampler = QwenImage21Resample(
out_dim, "downsample3d" if temporal_downsample else "downsample2d"
) if down_flag else None
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x.clone()
for resnet in self.resnets:
x = resnet(x)
if self.downsampler is not None:
x = self.downsampler(x)
return x + self.avg_shortcut(residual)
class QwenImage21Encoder3d(nn.Module):
def __init__(self, arch) -> None:
super().__init__()
dims = [arch.base_dim * factor for factor in [1, *arch.dim_mult]]
self.conv_in = QwenImage21CausalConv3d(arch.in_channels, dims[0], 3, padding=1)
self.down_blocks = nn.ModuleList()
scale = 1.0
for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
down_flag = index != len(arch.dim_mult) - 1
temporal = arch.temperal_downsample[index] if down_flag else False
if arch.is_residual:
self.down_blocks.append(QwenImage21ResidualDownBlock(
in_dim, out_dim, arch.dropout, arch.num_res_blocks, temporal, down_flag
))
else:
for _ in range(arch.num_res_blocks):
self.down_blocks.append(QwenImage21ResidualBlock(in_dim, out_dim, arch.dropout))
if scale in arch.attn_scales:
self.down_blocks.append(QwenImage21AttentionBlock(out_dim))
in_dim = out_dim
if down_flag:
mode = "downsample3d" if temporal else "downsample2d"
self.down_blocks.append(QwenImage21Resample(out_dim, mode))
scale /= 2.0
self.mid_block = QwenImage21MidBlock(out_dim, arch.dropout)
self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
self.conv_out = QwenImage21CausalConv3d(out_dim, arch.z_dim * 2, 3, padding=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.conv_in(x)
for block in self.down_blocks:
x = block(x)
x = self.mid_block(x)
return self.conv_out(F.silu(self.norm_out(x)))
class QwenImage21ResidualUpBlock(nn.Module):
def __init__(self, in_dim: int, out_dim: int, arch, temporal_upsample: bool, up_flag: bool) -> None:
super().__init__()
self.resnets = nn.ModuleList([
QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
for index in range(arch.num_res_blocks + 1)
])
self.avg_shortcut = QwenImage21DupUp3D(
in_dim, out_dim, 2 if temporal_upsample else 1, 2
) if up_flag else None
self.upsampler = QwenImage21Resample(
out_dim, "upsample3d" if temporal_upsample else "upsample2d", upsample_out_dim=out_dim
) if up_flag else None
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x.clone()
for resnet in self.resnets:
x = resnet(x)
if self.upsampler is not None:
x = self.upsampler(x)
if self.avg_shortcut is not None:
x = x + self.avg_shortcut(residual, first_chunk=True)
return x
class QwenImage21UpBlock(nn.Module):
def __init__(self, in_dim: int, out_dim: int, arch, upsample_mode: str | None) -> None:
super().__init__()
self.resnets = nn.ModuleList([
QwenImage21ResidualBlock(in_dim if index == 0 else out_dim, out_dim, arch.dropout)
for index in range(arch.num_res_blocks + 1)
])
self.upsamplers = nn.ModuleList([QwenImage21Resample(out_dim, upsample_mode)]) if upsample_mode else None
def forward(self, x: torch.Tensor) -> torch.Tensor:
for resnet in self.resnets:
x = resnet(x)
return self.upsamplers[0](x) if self.upsamplers is not None else x
class QwenImage21Decoder3d(nn.Module):
def __init__(self, arch) -> None:
super().__init__()
dim = arch.decoder_base_dim if arch.decoder_base_dim is not None else arch.base_dim
dims = [dim * factor for factor in [arch.dim_mult[-1], *arch.dim_mult[::-1]]]
temporal_upsample = arch.temperal_downsample[::-1]
self.conv_in = QwenImage21CausalConv3d(arch.z_dim, dims[0], 3, padding=1)
self.mid_block = QwenImage21MidBlock(dims[0], arch.dropout)
self.up_blocks = nn.ModuleList()
for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
up_flag = index != len(arch.dim_mult) - 1
temporal = temporal_upsample[index] if up_flag else False
if arch.is_residual:
self.up_blocks.append(QwenImage21ResidualUpBlock(in_dim, out_dim, arch, temporal, up_flag))
else:
if index > 0:
in_dim //= 2
mode = ("upsample3d" if temporal else "upsample2d") if up_flag else None
self.up_blocks.append(QwenImage21UpBlock(in_dim, out_dim, arch, mode))
self.norm_out = QwenImage21RMSNorm(out_dim, images=False)
self.conv_out = QwenImage21CausalConv3d(out_dim, arch.out_channels, 3, padding=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.mid_block(self.conv_in(x))
for block in self.up_blocks:
x = block(x)
return self.conv_out(F.silu(self.norm_out(x)))
def _patchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
if patch_size == 1:
return x
batch, channels, frames, height, width = x.shape
if height % patch_size or width % patch_size:
raise ValueError("Image height and width must be divisible by the VAE patch size.")
x = x.view(batch, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size)
x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous()
return x.view(batch, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size)
def _unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
if patch_size == 1:
return x
batch, channels, frames, height, width = x.shape
channels //= patch_size * patch_size
x = x.view(batch, channels, patch_size, patch_size, frames, height, width)
x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
return x.view(batch, channels, frames, height * patch_size, width * patch_size)
class AutoencoderKLQwenImage21(nn.Module):
_supports_gradient_checkpointing = False
_group_offload_block_modules = ["quant_conv", "post_quant_conv", "encoder", "decoder"]
def __init__(self, config: VAEConfig) -> None:
super().__init__()
self.config = config
if not config.load_encoder or not config.load_decoder:
raise ValueError("Qwen-Image-2.1 requires the complete encoder and decoder checkpoint.")
arch = config.arch_config
self.z_dim = arch.z_dim
self.latent_channels = arch.z_dim
self.encoder = QwenImage21Encoder3d(arch)
self.quant_conv = QwenImage21CausalConv3d(arch.z_dim * 2, arch.z_dim * 2, 1)
self.post_quant_conv = QwenImage21CausalConv3d(arch.z_dim, arch.z_dim, 1)
self.decoder = QwenImage21Decoder3d(arch)
self.spatial_compression_ratio = arch.scale_factor_spatial
self.temporal_compression_ratio = 1
self.use_slicing = False
self.use_tiling = config.use_tiling
self.tile_sample_min_height = config.tile_sample_min_height
self.tile_sample_min_width = config.tile_sample_min_width
self.tile_sample_stride_height = config.tile_sample_stride_height
self.tile_sample_stride_width = config.tile_sample_stride_width
self._validate_tiling()
@property
def device(self) -> torch.device:
return next(self.parameters()).device
@property
def dtype(self) -> torch.dtype:
return next(self.parameters()).dtype
def enable_slicing(self) -> None:
self.use_slicing = True
def disable_slicing(self) -> None:
self.use_slicing = False
def _validate_tiling(self) -> None:
ratio = self.spatial_compression_ratio
for extent, stride in ((self.tile_sample_min_height, self.tile_sample_stride_height),
(self.tile_sample_min_width, self.tile_sample_stride_width)):
if extent < ratio or stride < ratio or extent % ratio or stride % ratio or stride > extent:
raise ValueError(
"VAE tile size and stride must be positive spatial-scale multiples, with stride <= size."
)
def enable_tiling(self, tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None,
tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None,
**kwargs) -> None:
self.use_tiling = True
self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
self._validate_tiling()
def disable_tiling(self) -> None:
self.use_tiling = False
def clear_cache(self) -> None:
"""Image inference retains no temporal feature cache."""
@staticmethod
def _validate_image(x: torch.Tensor, channels: int, name: str) -> None:
if x.ndim != 5 or x.shape[1] != channels or x.shape[2] != 1:
raise ValueError(f"`{name}` must have shape [B, {channels}, 1, H, W], got {tuple(x.shape)}.")
if x.shape[0] < 1 or min(x.shape[-2:]) < 1:
raise ValueError(f"`{name}` must contain at least one nonempty image.")
def _encode(self, x: torch.Tensor) -> torch.Tensor:
if self.config.patch_size is not None:
x = _patchify(x, self.config.patch_size)
if self.use_tiling and (x.shape[-2] > self.tile_sample_min_height or x.shape[-1] > self.tile_sample_min_width):
return self.tiled_encode(x)
return self.quant_conv(self.encoder(x))
def encode(
self, x: torch.Tensor, return_dict: bool = True
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
self._validate_image(x, self.config.in_channels, "x")
if x.shape[-2] % self.spatial_compression_ratio or x.shape[-1] % self.spatial_compression_ratio:
raise ValueError("Image height and width must be divisible by the VAE spatial compression ratio.")
if self.use_slicing and x.shape[0] > 1:
moments = torch.cat([self._encode(item) for item in x.split(1)])
else:
moments = self._encode(x)
posterior = DiagonalGaussianDistribution(moments)
return AutoencoderKLOutput(posterior) if return_dict else (posterior,)
def _decode(self, z: torch.Tensor) -> torch.Tensor:
min_height = self.tile_sample_min_height // self.spatial_compression_ratio
min_width = self.tile_sample_min_width // self.spatial_compression_ratio
if self.use_tiling and (z.shape[-2] > min_height or z.shape[-1] > min_width):
return self.tiled_decode(z).sample
decoded = self.decoder(self.post_quant_conv(z))
if self.config.patch_size is not None:
decoded = _unpatchify(decoded, self.config.patch_size)
return decoded.clamp(-1.0, 1.0)
def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
self._validate_image(z, self.z_dim, "z")
if self.use_slicing and z.shape[0] > 1:
decoded = torch.cat([self._decode(item) for item in z.split(1)])
else:
decoded = self._decode(z)
return DecoderOutput(decoded) if return_dict else (decoded,)
@staticmethod
def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
for y in range(blend_extent):
b[:, :, :, y, :] = (
a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent)
+ b[:, :, :, y, :] * (y / blend_extent)
)
return b
@staticmethod
def blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
for x in range(blend_extent):
b[:, :, :, :, x] = (
a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent)
+ b[:, :, :, :, x] * (x / blend_extent)
)
return b
def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
height, width = x.shape[-2:]
ratio = self.spatial_compression_ratio // (self.config.patch_size or 1)
tile_height = self.tile_sample_min_height // ratio
tile_width = self.tile_sample_min_width // ratio
stride_height = self.tile_sample_stride_height // ratio
stride_width = self.tile_sample_stride_width // ratio
rows = []
for i in range(0, height, self.tile_sample_stride_height):
row = []
for j in range(0, width, self.tile_sample_stride_width):
tile = x[:, :, :, i:i + self.tile_sample_min_height, j:j + self.tile_sample_min_width]
row.append(self.quant_conv(self.encoder(tile)))
rows.append(row)
result_rows = []
for i, row in enumerate(rows):
result_row = []
for j, tile in enumerate(row):
if i > 0:
tile = self.blend_v(rows[i - 1][j], tile, tile_height - stride_height)
if j > 0:
tile = self.blend_h(row[j - 1], tile, tile_width - stride_width)
result_row.append(tile[:, :, :, :stride_height, :stride_width])
result_rows.append(torch.cat(result_row, dim=-1))
return torch.cat(result_rows, dim=3)[:, :, :, :height // ratio, :width // ratio]
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
height, width = z.shape[-2:]
ratio = self.spatial_compression_ratio
tile_height = self.tile_sample_min_height // ratio
tile_width = self.tile_sample_min_width // ratio
stride_height = self.tile_sample_stride_height // ratio
stride_width = self.tile_sample_stride_width // ratio
sample_height, sample_width = height * ratio, width * ratio
sample_stride_height = self.tile_sample_stride_height
sample_stride_width = self.tile_sample_stride_width
if self.config.patch_size is not None:
patch_size = self.config.patch_size
sample_height //= patch_size
sample_width //= patch_size
sample_stride_height //= patch_size
sample_stride_width //= patch_size
blend_height = self.tile_sample_min_height // patch_size - sample_stride_height
blend_width = self.tile_sample_min_width // patch_size - sample_stride_width
else:
blend_height = self.tile_sample_min_height - sample_stride_height
blend_width = self.tile_sample_min_width - sample_stride_width
rows = []
for i in range(0, height, stride_height):
row = []
for j in range(0, width, stride_width):
tile = z[:, :, :, i:i + tile_height, j:j + tile_width]
row.append(self.decoder(self.post_quant_conv(tile)))
rows.append(row)
result_rows = []
for i, row in enumerate(rows):
result_row = []
for j, tile in enumerate(row):
if i > 0:
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
if j > 0:
tile = self.blend_h(row[j - 1], tile, blend_width)
result_row.append(tile[:, :, :, :sample_stride_height, :sample_stride_width])
result_rows.append(torch.cat(result_row, dim=-1))
decoded = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
if self.config.patch_size is not None:
decoded = _unpatchify(decoded, self.config.patch_size)
decoded = decoded.clamp(-1.0, 1.0)
return DecoderOutput(decoded) if return_dict else (decoded,)
def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
return (latents - mean) / std
def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
mean = torch.tensor(self.config.latents_mean, device=latents.device, dtype=latents.dtype)
std = torch.tensor(self.config.latents_std, device=latents.device, dtype=latents.dtype)
mean, std = mean.view(1, self.z_dim, 1, 1, 1), std.view(1, self.z_dim, 1, 1, 1)
return latents * std + mean
def forward(self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True,
generator: torch.Generator | None = None) -> DecoderOutput | tuple[torch.Tensor]:
posterior = self.encode(sample).latent_dist
latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode()
return self.decode(latents, return_dict=return_dict)
@classmethod
def from_pretrained(cls, model_path: str | Path, *, subfolder: str | None = None,
config: VAEConfig | None = None, torch_dtype: torch.dtype | None = None,
device: torch.device | str = "cpu", local_files_only: bool = True) -> AutoencoderKLQwenImage21:
"""Strictly load a local Diffusers-layout VAE component without downloading weights."""
from safetensors.torch import load_file
if not local_files_only:
raise ValueError("Qwen-Image-2.1 VAE.from_pretrained accepts local checkpoints only.")
path = Path(model_path)
if subfolder:
path /= subfolder
if not path.is_dir():
raise FileNotFoundError(f"Local Qwen-Image-2.1 VAE directory does not exist: {path}")
with (path / "config.json").open() as handle:
arch_fields = json.load(handle)
class_name = arch_fields.get("_class_name")
if class_name != "AutoencoderKLQwenImage21":
raise ValueError(f"Expected AutoencoderKLQwenImage21 checkpoint, got {class_name!r}.")
config = config or QwenImage21VAEConfig()
config.update_model_arch(arch_fields)
model = cls(config)
files = sorted(path.glob("*.safetensors"))
if not files:
raise FileNotFoundError(f"No safetensors VAE checkpoint found in {path}.")
state = {}
for file in files:
shard = load_file(str(file))
duplicates = state.keys() & shard.keys()
if duplicates:
raise ValueError(f"Duplicate tensors in VAE checkpoint shards: {sorted(duplicates)[:5]}")
state.update(shard)
model.load_state_dict(state, strict=True)
logger.info("Loaded Qwen-Image-2.1 VAE strictly from %s", path)
return model.to(device=device, dtype=torch_dtype).eval()
EntryClass = AutoencoderKLQwenImage21
fastvideo/pipelines/basic/qwen_image21/__init__.py
+1
View File
@@ -0,0 +1 @@
# SPDX-License-Identifier: Apache-2.0
fastvideo/pipelines/basic/qwen_image21/presets.py
+39
View File
@@ -0,0 +1,39 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 defaults shared by text and reference image generation."""
from fastvideo.api.presets import InferencePreset, PresetStageSpec
QWEN_IMAGE21 = InferencePreset(
name="qwen_image21",
version=1,
model_family="qwen_image21",
description="Qwen-Image-2.1 image generation, editing and ordered references",
workload_type="t2i",
stage_schemas=(PresetStageSpec(name="denoise",
kind="denoising",
description="Qwen-Image-2.1 Euler flow steps",
allowed_overrides=frozenset({"num_inference_steps", "true_cfg_scale"})), ),
defaults={
"height": 1024,
"width": 1024,
"num_frames": 1,
"num_videos_per_prompt": 1,
"fps": 1,
"seed": 42,
"num_inference_steps": 40,
"guidance_scale": 1.0,
"true_cfg_scale": 1.0,
},
)
QWEN_IMAGE21_EDIT = InferencePreset(
name="qwen_image21_edit",
version=1,
model_family="qwen_image21",
description="Qwen-Image-2.1 image edits and reference-to-image generation",
workload_type="i2i",
stage_schemas=QWEN_IMAGE21.stage_schemas,
defaults=dict(QWEN_IMAGE21.defaults),
)
ALL_PRESETS = (QWEN_IMAGE21, QWEN_IMAGE21_EDIT)
fastvideo/pipelines/basic/qwen_image21/qwen_image21_pipeline.py
+36
View File
@@ -0,0 +1,36 @@
# SPDX-License-Identifier: Apache-2.0
"""Native Qwen-Image-2.1 pipeline for generation and image-conditioned edits."""
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
from .stages import (
QwenImage21DecodingStage,
QwenImage21DenoisingStage,
QwenImage21InputStage,
QwenImage21LatentStage,
QwenImage21ScheduleStage,
QwenImage21TextEncodingStage,
)
class QwenImage21Pipeline(ComposedPipelineBase):
pipeline_config_cls = QwenImage21PipelineConfig
sampling_params_cls = SamplingParam
_required_config_modules = ["processor", "text_encoder", "transformer", "vae", "scheduler"]
_lazy_module_names = ("text_encoder", "transformer", "vae")
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
self.add_stage("input_preparation", QwenImage21InputStage())
self.add_stage("text_encoding",
QwenImage21TextEncodingStage(self.get_module("text_encoder"), self.get_module("processor")))
self.add_stage("latent_preparation", QwenImage21LatentStage(self.get_module("vae")))
self.add_stage("timestep_preparation", QwenImage21ScheduleStage(self.get_module("scheduler")))
self.add_stage("denoising",
QwenImage21DenoisingStage(self.get_module("transformer"), self.get_module("scheduler")))
self.add_stage("decoding", QwenImage21DecodingStage(self.get_module("vae")))
EntryClass = QwenImage21Pipeline
fastvideo/pipelines/basic/qwen_image21/stages.py
+510
View File
@@ -0,0 +1,510 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved.
# Adapted from Diffusers' Apache-2.0 QwenImage21Pipeline.
"""Separate preparation, encoding, scheduling, denoising and decoding stages."""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import numpy as np
import torch
from PIL import Image, ImageOps
from torch.distributed.tensor import DTensor
from tqdm.auto import tqdm
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.hooks.hooks import ModuleHookManager
from fastvideo.logger import init_logger
from fastvideo.models import pinned_offload
from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache, QwenImage21KVCacheAllocationError
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import V, VerificationResult
logger = init_logger(__name__)
STATE_KEY = "qwen_image21"
SYSTEM_PROMPT = "Comprehend and analyze the provided prompt."
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
@dataclass
class QwenImage21State:
references: list[Image.Image] = field(default_factory=list)
img_shapes: list[list[tuple[int, int, int]]] = field(default_factory=list)
image_pad_mask: torch.Tensor | None = None
negative_image_pad_mask: torch.Tensor | None = None
def load_reference(value: Any) -> Image.Image:
"""Detach a local image from its file and preserve alpha for the VAE."""
if isinstance(value, str | Path):
with Image.open(value) as image:
return ImageOps.exif_transpose(image).convert("RGBA")
if isinstance(value, Image.Image):
return ImageOps.exif_transpose(value).convert("RGBA")
if isinstance(value, np.ndarray) and value.dtype == np.uint8 and value.ndim in (2, 3):
return Image.fromarray(value).convert("RGBA")
raise TypeError("Each reference must be a local image path, PIL image, or uint8 image array")
def collect_references(references: list[Any] | None, image_path: Any, pil_image: Any) -> list[Image.Image]:
if references is not None and not isinstance(references, list | tuple):
raise TypeError("references must be an ordered list of images")
shortcuts = [value for value in (image_path, pil_image) if value is not None]
if len(shortcuts) > 1 or (references and shortcuts):
raise ValueError("Supply references, image_path, or pil_image as a single source of reference images")
values = list(references or shortcuts)
if len(values) > 10:
raise ValueError("Qwen-Image-2.1 supports at most 10 ordered reference images")
return [load_reference(value) for value in values]
def reference_dimensions(width: int, height: int, resolution: int) -> tuple[int, int]:
if min(width, height, resolution) <= 0:
raise ValueError("Image dimensions and reference_resolution must be positive")
ratio = width / height
width_float = math.sqrt(resolution**2 * ratio)
height_float = width_float / ratio
new_width = round(width_float / 32) * 32
new_height = round(height_float / 32) * 32
if min(new_width, new_height) < 32:
raise ValueError("Reference aspect ratio is too extreme for a 32-pixel latent grid")
return new_width, new_height
def resize_reference(image: Image.Image, resolution: int) -> Image.Image:
return image.resize(reference_dimensions(image.width, image.height, resolution), resample=Image.Resampling.LANCZOS)
def vision_reference(image: Image.Image) -> Image.Image:
white = Image.new("RGB", image.size, (255, 255, 255))
white.paste(image, mask=image.getchannel("A"))
return white
def reference_pixels(image: Image.Image) -> torch.Tensor:
# [B, RGBA, one frame, H, W], with all four channels normalized to [-1, 1].
# Build the batch axis before permuting, as upstream's numpy_to_pt does.
# Even a singleton batch stride affects BF16 convolution kernel rounding.
pixels = torch.from_numpy(np.array(image, dtype=np.float32)[None]).permute(0, 3, 1, 2)
return (pixels / 127.5 - 1).unsqueeze(2)
def prompt_template(prompt: str, num_references: int) -> str:
if not 0 <= num_references <= 10:
raise ValueError("Qwen-Image-2.1 supports 0–10 references")
prefix = " ".join(f"<image{index + 1}>{IMAGE_PLACEHOLDER}" for index in range(num_references))
return (f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
f"<|im_start|>user\n{prefix}{prompt or ' '}<|im_end|>\n<|im_start|>assistant\n")
def extract_prompt(hidden: torch.Tensor, input_ids: torch.Tensor, attention_mask: torch.Tensor, drop_tokens: int,
image_token_id: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
rows, image_rows = [], []
for embeddings, ids, mask in zip(hidden, input_ids, attention_mask, strict=True):
valid = mask.bool()
rows.append(embeddings[valid][drop_tokens:])
image_rows.append((ids[valid] == image_token_id)[drop_tokens:])
if not rows or min(row.shape[0] for row in rows) == 0:
raise ValueError("The processor produced an empty prompt after removing system tokens")
length = max(row.shape[0] for row in rows)
embeds = torch.stack([torch.cat((row, row.new_zeros(length - row.shape[0], row.shape[1]))) for row in rows])
masks = torch.stack([torch.arange(length, device=row.device) < row.shape[0] for row in rows])
image_masks = torch.stack([torch.cat((row, row.new_zeros(length - row.shape[0]))) for row in image_rows])
return embeds, masks, image_masks
def pack_latents(latents: torch.Tensor) -> torch.Tensor:
"""Pack [B, C, 1, H, W] into [B, H*W, C], with width varying fastest."""
if latents.ndim != 5 or latents.shape[2] != 1:
raise ValueError("Qwen-Image-2.1 latents must have shape [B, C, 1, H, W]")
return latents.flatten(2).transpose(1, 2)
def unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tensor:
"""Restore [B, C, 1, H/16, W/16] from tokens containing all C channels."""
if latents.ndim != 3 or latents.shape[1] != (height // 16) * (width // 16):
raise ValueError("Packed latent length does not match the requested image dimensions")
return latents.transpose(1, 2).reshape(latents.shape[0], latents.shape[2], 1, height // 16, width // 16)
def normalize_latents(latents: torch.Tensor, means: Any, stds: Any, *, inverse: bool = False) -> torch.Tensor:
mean = torch.as_tensor(means, device=latents.device, dtype=latents.dtype).reshape(1, -1, 1, 1, 1)
std = torch.as_tensor(stds, device=latents.device, dtype=latents.dtype).reshape(1, -1, 1, 1, 1)
if mean.shape[1] != latents.shape[1] or std.shape[1] != latents.shape[1] or bool((std <= 0).any()):
raise ValueError("VAE latent statistics must match its channels and have positive standard deviations")
return latents * std + mean if inverse else (latents - mean) / std
def schedule_shift(target_tokens: int, config: Any) -> float:
base_len, max_len = config.get("base_image_seq_len", 256), config.get("max_image_seq_len", 8192)
base_shift, max_shift = config.get("base_shift", 0.5), config.get("max_shift", 0.9)
slope = (max_shift - base_shift) / (max_len - base_len)
return base_shift + slope * (target_tokens - base_len)
def validate_request(batch: Any, args: Any) -> None:
if args.tp_size != 1 or args.sp_size != 1 or args.num_gpus != 1:
raise ValueError("The Qwen-Image-2.1 port currently supports one GPU with tp_size=sp_size=1")
if batch.num_frames != 1 or batch.num_videos_per_prompt != 1:
raise ValueError("Qwen-Image-2.1 requires num_frames=1 and num_videos_per_prompt=1")
if not isinstance(batch.prompt, str) and not (isinstance(batch.prompt, list) and len(batch.prompt) == 1
and isinstance(batch.prompt[0], str)):
raise ValueError("Qwen-Image-2.1 requires one text prompt per request")
for name, value in (("height", batch.height), ("width", batch.width), ("reference_resolution",
args.pipeline_config.reference_resolution)):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0 or value % 32:
raise ValueError(f"{name} must be a positive multiple of 32")
if not isinstance(batch.num_inference_steps, int) or batch.num_inference_steps < 1:
raise ValueError("num_inference_steps must be positive")
if not math.isfinite(batch.true_cfg_scale) or batch.true_cfg_scale < 1:
raise ValueError("true_cfg_scale must be finite and at least 1")
if batch.guidance_scale != 1:
raise ValueError("Qwen-Image-2.1 has no embedded guidance; use true_cfg_scale with negative_prompt")
if batch.video_path is not None or getattr(batch, "conditioning_mask", None) is not None:
raise ValueError("Pass annotations or mask images as references; this pipeline takes single-frame images")
if batch.prompt_embeds or batch.negative_prompt_embeds:
raise ValueError("Externally supplied prompt embeddings are not supported by this Qwen-Image-2.1 port")
if batch.negative_prompt is not None and not isinstance(batch.negative_prompt, str):
raise ValueError("negative_prompt must be a string or None")
def _state(batch: ForwardBatch) -> QwenImage21State:
state = batch.extra.get(STATE_KEY)
if not isinstance(state, QwenImage21State):
raise ValueError("Qwen-Image-2.1 input preparation must precede this stage")
return state
class QwenImage21InputStage(PipelineStage):
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
result = VerificationResult()
result.add_check("height", batch.height, V.positive_int_divisible(32))
result.add_check("width", batch.width, V.positive_int_divisible(32))
result.add_check("num_frames", batch.num_frames, lambda value: value == 1)
return result
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check(STATE_KEY, batch.extra.get(STATE_KEY),
lambda value: isinstance(value, QwenImage21State))
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
validate_request(batch, fastvideo_args)
references = collect_references(batch.references, batch.image_path, batch.pil_image)
references = [
resize_reference(image, fastvideo_args.pipeline_config.reference_resolution) for image in references
]
shapes = [(1, image.height // 16, image.width // 16) for image in references]
shapes.append((1, batch.height // 16, batch.width // 16))
batch.extra[STATE_KEY] = QwenImage21State(references=references, img_shapes=[shapes])
batch.batch_size = 1
batch.do_classifier_free_guidance = batch.true_cfg_scale > 1 and batch.negative_prompt is not None
if batch.true_cfg_scale > 1 and not batch.do_classifier_free_guidance:
logger.warning("true_cfg_scale > 1 requires an explicit negative_prompt; guidance is disabled")
return batch
class QwenImage21TextEncodingStage(PipelineStage):
performance_component_metric = "text_encoder_time_s"
def __init__(self, text_encoder, processor) -> None:
self.text_encoder = text_encoder
self.processor = processor
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check(STATE_KEY, batch.extra.get(STATE_KEY),
lambda value: isinstance(value, QwenImage21State))
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("prompt_embeds", batch.prompt_embeds,
lambda value: len(value) == 1 and value[0].ndim == 3)
def _encode(self, prompt: str, state: QwenImage21State, device: torch.device):
kwargs = dict(text=[prompt_template(prompt, len(state.references))],
padding=True,
padding_side="left",
return_tensors="pt")
if state.references:
kwargs["images"] = [vision_reference(image) for image in state.references]
model_inputs = self.processor(**kwargs)
inputs = {key: value.to(device) for key, value in model_inputs.items() if isinstance(value, torch.Tensor)}
hidden = self.text_encoder(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
pixel_values=inputs.get("pixel_values"),
image_grid_thw=inputs.get("image_grid_thw"),
mm_token_type_ids=inputs.get("mm_token_type_ids"),
)
system = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
system_ids = self.processor.apply_chat_template(system, tokenize=True, return_dict=False)
drop_tokens = len(system_ids[0])
image_token_id = self.processor.tokenizer.encode("<|image_pad|>")[0]
embeds, mask, image_mask = extract_prompt(hidden, inputs["input_ids"], inputs["attention_mask"], drop_tokens,
image_token_id)
return embeds, None if bool(mask.all()) else mask, image_mask
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
state = _state(batch)
device = get_local_torch_device()
first_param = next(self.text_encoder.parameters(), None)
# The loader's FSDP2 CPU offload already streams individual text layers.
move = (fastvideo_args.text_encoder_cpu_offload and first_param is not None
and not isinstance(first_param, DTensor))
try:
if move:
pinned_offload.load(self.text_encoder, device, pin=fastvideo_args.pin_cpu_memory)
prompt = batch.prompt[0] if isinstance(batch.prompt, list) else batch.prompt
if prompt is None:
raise ValueError("A text prompt is required")
embeds, mask, state.image_pad_mask = self._encode(prompt, state, device)
batch.prompt_embeds = [embeds]
batch.prompt_attention_mask = [mask]
if batch.do_classifier_free_guidance:
if batch.negative_prompt is None:
raise ValueError("True CFG requires an explicit negative prompt")
embeds, mask, state.negative_image_pad_mask = self._encode(batch.negative_prompt, state, device)
batch.negative_prompt_embeds = [embeds]
batch.negative_attention_mask = [mask]
finally:
if move:
pinned_offload.unload(self.text_encoder)
return batch
class QwenImage21LatentStage(PipelineStage):
performance_component_metric = "vae_encode_time_s"
def __init__(self, vae) -> None:
self.vae = vae
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("prompt_embeds", batch.prompt_embeds, lambda value: len(value) == 1)
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
state = _state(batch)
device, dtype = get_local_torch_device(), batch.prompt_embeds[0].dtype
channels = fastvideo_args.pipeline_config.dit_config.arch_config.in_channels
shape = (1, channels, 1, batch.height // 16, batch.width // 16)
if batch.latents is None:
generator = batch.generator
if isinstance(generator, list):
if len(generator) != 1:
raise ValueError("Qwen-Image-2.1 requires one generator")
generator = generator[0]
if generator is None:
generator = torch.Generator(device="cpu").manual_seed(batch.seed or 0)
noise_device = generator.device
batch.latents = pack_latents(
torch.randn(shape, generator=generator, device=noise_device, dtype=dtype).to(device))
else:
if batch.latents.ndim == 5:
batch.latents = pack_latents(batch.latents)
if tuple(batch.latents.shape) != (1, shape[3] * shape[4], channels):
raise ValueError("Supplied latents must match [1, target tokens, latent channels]")
batch.latents = batch.latents.to(device=device, dtype=dtype)
batch.raw_latent_shape = shape
if state.references:
try:
pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory)
if fastvideo_args.pipeline_config.vae_tiling:
self.vae.enable_tiling()
else:
self.vae.disable_tiling()
vae_dtype = next(self.vae.parameters()).dtype
rows = []
for image in state.references:
pixels = reference_pixels(image).to(device=device, dtype=vae_dtype)
latents = self.vae.encode(pixels).latent_dist.mode()
latents = normalize_latents(latents, self.vae.config.latents_mean, self.vae.config.latents_std)
rows.append(pack_latents(latents).to(dtype))
batch.image_latent = torch.cat(rows, dim=1)
expected = int(state.image_pad_mask.sum().item()) * 4
if expected != batch.image_latent.shape[1]:
raise ValueError("The processor's image slots do not match the VAE reference grid. "
"Use the checkpoint's processor and a reference_resolution "
"within its image budget.")
finally:
if fastvideo_args.vae_cpu_offload:
pinned_offload.unload(self.vae)
# One target mask slot stands for four unpatched latent tokens.
slots = state.image_pad_mask.new_ones(1, batch.latents.shape[1] // 4)
state.image_pad_mask = torch.cat((state.image_pad_mask, slots), dim=1)
if state.negative_image_pad_mask is not None:
slots = state.negative_image_pad_mask.new_ones(1, batch.latents.shape[1] // 4)
state.negative_image_pad_mask = torch.cat((state.negative_image_pad_mask, slots), dim=1)
return batch
class QwenImage21ScheduleStage(PipelineStage):
def __init__(self, scheduler) -> None:
self.scheduler = scheduler
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("timesteps", batch.timesteps, V.with_dims(1))
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
sigmas = batch.sigmas.tolist() if isinstance(batch.sigmas, torch.Tensor) else batch.sigmas
if sigmas is None:
sigmas = np.linspace(1.0, 1 / batch.num_inference_steps, batch.num_inference_steps).tolist()
if not sigmas or any(not np.isfinite(sigma) or sigma <= 0 or sigma > 1 for sigma in sigmas):
raise ValueError("sigmas must be finite values in (0, 1]")
if any(first <= second for first, second in zip(sigmas, sigmas[1:], strict=False)):
raise ValueError("sigmas must be strictly decreasing")
self.scheduler.set_timesteps(sigmas=sigmas,
device=get_local_torch_device(),
mu=schedule_shift(batch.latents.shape[1], self.scheduler.config))
self.scheduler.set_begin_index(0)
batch.timesteps = self.scheduler.timesteps
batch.num_inference_steps = len(batch.timesteps)
return batch
class QwenImage21DenoisingStage(PipelineStage):
performance_component_metric = "dit_time_s"
def __init__(self, transformer, scheduler) -> None:
self.transformer, self.scheduler = transformer, scheduler
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("latents", batch.latents,
V.with_dims(3)).add_check("timesteps", batch.timesteps, V.with_dims(1))
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
state, device = _state(batch), get_local_torch_device()
cache_device = fastvideo_args.pipeline_config.kv_cache_device
if cache_device not in ("cpu", "cuda"):
raise ValueError("kv_cache_device must be cpu or cuda; use_kv_cache=False disables caching")
cache_enabled = fastvideo_args.pipeline_config.use_kv_cache and self.transformer.config.causal_condition
caches = []
if cache_enabled:
caches = [QwenImage21KVCache(len(self.transformer.transformer_blocks), storage_device=cache_device)]
if batch.do_classifier_free_guidance:
caches.append(QwenImage21KVCache(len(self.transformer.transformer_blocks), storage_device=cache_device))
move = (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
and not fastvideo_args.use_fsdp_inference)
try:
if move:
pinned_offload.load(self.transformer, device, pin=fastvideo_args.pin_cpu_memory)
for index, timestep in enumerate(tqdm(batch.timesteps, desc="Denoising", total=len(batch.timesteps))):
if getattr(self, "interrupt", False):
break
latent_input = batch.latents
if batch.image_latent is not None:
latent_input = torch.cat((batch.image_latent, latent_input), dim=1)
mode = "extract" if cache_enabled and index == 0 else ("cached" if cache_enabled else None)
def predict(embeds: torch.Tensor, mask: torch.Tensor | None, image_mask: torch.Tensor,
cache: QwenImage21KVCache | None, latent_input: torch.Tensor, timestep: torch.Tensor,
mode: str | None) -> torch.Tensor:
result = self.transformer(
hidden_states=latent_input,
timestep=timestep.expand(1).to(batch.latents.dtype) / 1000,
encoder_hidden_states=embeds,
encoder_hidden_states_mask=mask,
img_shapes=state.img_shapes,
img_mask=image_mask,
kv_cache=cache,
kv_cache_mode=mode,
return_dict=False,
)[0]
return result[:, -batch.latents.shape[1]:]
while True:
try:
positive = predict(batch.prompt_embeds[0], batch.prompt_attention_mask[0], state.image_pad_mask,
caches[0] if caches else None, latent_input, timestep, mode)
noise = positive
if batch.do_classifier_free_guidance:
negative = predict(batch.negative_prompt_embeds[0], batch.negative_attention_mask[0],
state.negative_image_pad_mask, caches[1] if caches else None,
latent_input, timestep, mode)
noise = negative + batch.true_cfg_scale * (positive - negative)
break
except QwenImage21KVCacheAllocationError:
if not cache_enabled:
raise
logger.warning("Host prefix-cache allocation failed; retrying this request without KV caching")
for cache in caches:
cache.clear()
caches = []
cache_enabled, mode = False, None
# An aborted layer may retain the next layer's prefetch.
self._release_layerwise()
batch.latents = self.scheduler.step(noise, timestep, batch.latents, return_dict=False)[0]
batch.step_index, batch.timestep = index, timestep
finally:
for cache in caches:
cache.clear()
if fastvideo_args.dit_layerwise_offload:
self._release_layerwise()
if move:
pinned_offload.unload(self.transformer)
return batch
def _release_layerwise(self) -> None:
# The shared offloader links the last block back to the first, leaving
# its next-step prefetch resident. Also handle an interrupted block.
for block in self.transformer.transformer_blocks:
manager = ModuleHookManager.get_from(block)
hook = manager.get_forward_hook("LayerwiseOffloadHook") if manager is not None else None
if hook is not None and hook.state.gpu_named_parameters:
torch.cuda.current_stream().wait_stream(hook.state.async_copy_stream)
hook.state.release_gpu_params()
class QwenImage21DecodingStage(PipelineStage):
performance_component_metric = "vae_decode_time_s"
def __init__(self, vae) -> None:
self.vae = vae
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("latents", batch.latents, V.with_dims(3))
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
return VerificationResult().add_check("output", batch.output, V.is_tensor)
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if fastvideo_args.output_type == "latent":
batch.output = batch.latents
return batch
device = get_local_torch_device()
try:
pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory)
if fastvideo_args.pipeline_config.vae_tiling:
self.vae.enable_tiling()
else:
self.vae.disable_tiling()
dtype = next(self.vae.parameters()).dtype
latents = unpack_latents(batch.latents, batch.height, batch.width).to(device=device, dtype=dtype)
latents = normalize_latents(latents,
self.vae.config.latents_mean,
self.vae.config.latents_std,
inverse=True)
# Keep alpha as the fourth channel all the way to the PNG exporter.
batch.output = self.vae.decode(latents, return_dict=False)[0].div(2).add(0.5).clamp(0, 1).float()
finally:
if fastvideo_args.vae_cpu_offload:
pinned_offload.unload(self.vae)
return batch
fastvideo/registry.py
+17
View File
@@ -55,6 +55,7 @@ from fastvideo.models.wan.definition import WAN_MODEL_DEFINITIONS, WanModelDefin
from fastvideo.configs.pipelines.flux import FluxPipelineConfig
from fastvideo.configs.pipelines.sd35 import SD35Config
from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
@@ -1204,6 +1205,19 @@ def _register_configs() -> None:
],
)
register_configs(
sampling_param_cls=None,
pipeline_config_cls=QwenImage21PipelineConfig,
workload_types=(WorkloadType.T2I, WorkloadType.I2I),
hf_model_paths=["Qwen/Qwen-Image-2.1"],
model_detectors=[
lambda path: "qwenimage21pipeline" in path or "qwen-image-2.1" in path or "qwen_image21" in path,
],
model_family="qwen_image21",
default_preset="qwen_image21",
pipeline_cls_name="QwenImage21Pipeline",
)
# --- Part 3: Main Resolver ---
@@ -1337,6 +1351,8 @@ def _register_presets() -> None:
ALL_PRESETS as STABLE_AUDIO_PRESETS, )
from fastvideo.pipelines.basic.wan.presets import (
ALL_PRESETS as WAN_PRESETS, )
from fastvideo.pipelines.basic.qwen_image21.presets import (
ALL_PRESETS as QWEN_IMAGE21_PRESETS, )
from fastvideo.pipelines.basic.flux_2.presets import (
ALL_PRESETS as FLUX2_PRESETS, )
@@ -1364,6 +1380,7 @@ def _register_presets() -> None:
SD35_PRESETS,
STABLE_AUDIO_PRESETS,
WAN_PRESETS,
QWEN_IMAGE21_PRESETS,
)
for group in all_preset_groups:
for preset in group:
fastvideo/tests/golden_gate/test_qwen_image21.py
+102
View File
@@ -0,0 +1,102 @@
# SPDX-License-Identifier: Apache-2.0
"""Golden-gate: Qwen-Image-2.1 transformer block 0 with real BF16 weights.
The small packed sequence contains text, one reference image, more text, and
the target image. Causal text segments use masked SDPA; reference and target
segments use FlashAttention, including their unequal query/key lengths. Complex RoPE and
the two modulation rows match the causal-conditioning block interface.
Requires a BF16-capable CUDA GPU and FlashAttention. Select FA4 with
FASTVIDEO_FA4=1 before starting Python. The shared harness downloads
only checkpoint shards containing block 0 and compares a device/runtime-matched
golden exactly; it never loads the text encoder or VAE.
"""
from __future__ import annotations
import torch
from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate
__all__ = ["distributed_runtime"]
DIM = 4096 # 32 heads x 128, Qwen-Image-2.1 transformer/config.json
TEXT_BEFORE, TEXT_AFTER = 32, 16
REFERENCE_H, REFERENCE_W = 4, 4
TARGET_H, TARGET_W = 8, 16
CHECKPOINT_REVISION = "d26bb61231c349cf6b7896fa83353113880e1ba3"
def _build_block(layer: int) -> torch.nn.Module:
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig
from fastvideo.models.dits.qwen_image21 import QwenImage21TransformerBlock
from fastvideo.platforms import AttentionBackendEnum
arch = QwenImage21ArchConfig()
# The selector reads the construction dtype before the harness casts the
# loaded block. Keep it BF16 here so FLASH_ATTN cannot fall back to SDPA.
previous_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.bfloat16)
try:
block = QwenImage21TransformerBlock(
dim=arch.num_attention_heads * arch.attention_head_dim,
num_attention_heads=arch.num_attention_heads,
attention_head_dim=arch.attention_head_dim,
mlp_ratio=arch.mlp_ratio,
eps=arch.eps,
requested_backend=AttentionBackendEnum.FLASH_ATTN,
supported_attention_backends=arch._supported_attention_backends,
)
finally:
torch.set_default_dtype(previous_dtype)
assert block.attn.dense_attn is not None, "Qwen golden requires dense FlashAttention"
assert block.attn.dense_attn.backend is AttentionBackendEnum.FLASH_ATTN
return block
def _make_inputs(device: torch.device, seed: int) -> dict:
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig
from fastvideo.models.dits.qwen_image21 import (
QwenImage21Rope,
QwenImage21Transformer2DModel,
_qwenimage21_prefix_segments,
)
generator = torch.Generator(device="cpu").manual_seed(seed)
reference_tokens = REFERENCE_H * REFERENCE_W
target_tokens = TARGET_H * TARGET_W
image_pad_mask = torch.tensor(
[False] * TEXT_BEFORE + [True] * reference_tokens + [False] * TEXT_AFTER + [True] * target_tokens,
dtype=torch.bool,
device=device,
)
img_shapes = [(1, REFERENCE_H, REFERENCE_W), (1, TARGET_H, TARGET_W)]
image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes)
rope = QwenImage21Rope(theta=10000, axes_dim=QwenImage21ArchConfig().axes_dims_rope)
def rnd(*shape):
return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device=device, dtype=torch.bfloat16)
return {
"hidden_states": rnd(1, image_pad_mask.numel(), DIM),
# Row 0 conditions target tokens; row 1 is the prefix's zero-timestep row.
"modulation": rnd(2, 4 * DIM),
"target_token_mask": target_token_mask,
"rotary_emb": rope(img_shapes, image_pad_mask, device),
"segments": _qwenimage21_prefix_segments(image_ids, image_pad_mask.numel() - target_tokens),
}
SPEC = GateSpec(
name="qwen_image21",
repo_id="Qwen/Qwen-Image-2.1",
revision=CHECKPOINT_REVISION,
build_block=_build_block,
make_inputs=_make_inputs,
prefix_template="transformer_blocks.{N}.",
attention_backend="FLASH_ATTN",
)
def test_qwen_image21_golden_gate(distributed_runtime) -> None:
run_gate(SPEC)
fastvideo/tests/ssim/test_qwen_image21_similarity.py
+143
View File
@@ -0,0 +1,143 @@
# SPDX-License-Identifier: Apache-2.0
"""Single-GPU Qwen T2I and editing regressions, including exported alpha.
Requires the Qwen checkpoint, FlashAttention, and device-matched PNG references.
The shared SSIM harness compares all four RGBA channels and records metrics.
"""
from __future__ import annotations
from pathlib import Path
import pytest
import torch
from PIL import Image, ImageDraw
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.logger import init_logger
from fastvideo.tests.ssim.inference_similarity_utils import (
DEVICE_MAPPINGS,
run_text_to_video_similarity_test,
)
from fastvideo.tests.ssim.reference_utils import (
build_generated_output_dir,
get_cuda_device_name,
resolve_device_reference_folder,
select_ssim_params,
)
logger = init_logger(__name__)
REQUIRED_GPUS = 1
MODEL_ID = "Qwen/Qwen-Image-2.1"
MODEL_REVISION = "d26bb61231c349cf6b7896fa83353113880e1ba3"
MIN_ACCEPTABLE_SSIM = 0.98
QWEN_IMAGE21_DEFAULT_PARAMS: dict[str, object] = {
"model_path": MODEL_ID,
"num_gpus": REQUIRED_GPUS,
"sp_size": 1,
"tp_size": 1,
"height": 512,
"width": 512,
"num_frames": 1,
"fps": 1,
"num_inference_steps": 4,
"guidance_scale": 1.0,
"seed": 21,
}
QWEN_IMAGE21_FULL_QUALITY_PARAMS: dict[str, object] = {
**QWEN_IMAGE21_DEFAULT_PARAMS,
"height": 1024,
"width": 1024,
"num_inference_steps": 40,
}
QWEN_IMAGE21_MODEL_TO_PARAMS = {"Qwen/Qwen-Image-2.1": QWEN_IMAGE21_DEFAULT_PARAMS}
QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS = {"Qwen/Qwen-Image-2.1": QWEN_IMAGE21_FULL_QUALITY_PARAMS}
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Qwen image SSIM requires one CUDA GPU")
def _run_similarity(model_id: str, attention_backend_name: str, prompt: str,
references: list[str] | None = None) -> None:
device_reference_folder = resolve_device_reference_folder(
DEVICE_MAPPINGS + (("RTX 5090", "RTX5090"), ("RTX 4090", "RTX4090")),
device_name=get_cuda_device_name(),
logger=logger,
)
if device_reference_folder is None:
pytest.skip(f"No Qwen image SSIM reference mapping for {get_cuda_device_name()}")
script_dir = str(Path(__file__).resolve().parent)
selected_params = select_ssim_params(QWEN_IMAGE21_MODEL_TO_PARAMS,
QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS)[model_id]
run_text_to_video_similarity_test(
logger=logger,
script_dir=script_dir,
device_reference_folder=device_reference_folder,
prompt=prompt,
attention_backend_name=attention_backend_name,
model_id=model_id,
default_params_map=QWEN_IMAGE21_MODEL_TO_PARAMS,
full_quality_params_map=QWEN_IMAGE21_FULL_QUALITY_MODEL_TO_PARAMS,
min_acceptable_ssim=MIN_ACCEPTABLE_SSIM,
media_extension=".png",
init_kwargs_override={
"pipeline_config": QwenImage21PipelineConfig(
reference_resolution=selected_params["height"],
use_kv_cache=True,
),
"revision": MODEL_REVISION,
"workload_type": "i2i" if references else "t2i",
"use_fsdp_inference": False,
"dit_cpu_offload": True,
"dit_layerwise_offload": True,
"text_encoder_cpu_offload": True,
"vae_cpu_offload": True,
"lazy_module_load": True,
"pin_cpu_memory": False,
"vae_tiling": True,
},
generation_kwargs_override={
"references": references,
"negative_prompt": None,
"true_cfg_scale": 1.0,
"return_frames": True,
"save_video": True,
},
)
output_dir = build_generated_output_dir(script_dir, device_reference_folder, model_id, attention_backend_name)
with Image.open(Path(output_dir) / f"{prompt[:100].strip()}.png") as generated:
assert generated.mode == "RGBA", "Qwen PNG export must retain its alpha channel"
assert generated.size == (selected_params["width"], selected_params["height"])
@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
@pytest.mark.parametrize("model_id", list(QWEN_IMAGE21_MODEL_TO_PARAMS))
def test_qwen_image21_t2i_similarity(model_id: str, attention_backend_name: str) -> None:
_run_similarity(
model_id,
attention_backend_name,
"A studio photograph of a blue ceramic vase on a white background",
)
@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
@pytest.mark.parametrize("model_id", list(QWEN_IMAGE21_MODEL_TO_PARAMS))
def test_qwen_image21_edit_similarity(model_id: str, attention_backend_name: str, tmp_path: Path) -> None:
reference = Image.new("RGB", (512, 512), "white")
draw = ImageDraw.Draw(reference)
draw.ellipse((176, 160, 336, 432), fill=(32, 96, 224))
draw.rectangle((224, 96, 288, 192), fill=(32, 96, 224))
draw.ellipse((224, 88, 288, 112), fill=(32, 96, 224))
reference_path = tmp_path / "blue-vase.png"
reference.save(reference_path)
_run_similarity(
model_id,
attention_backend_name,
"Change the blue vase to red and keep the white background",
references=[str(reference_path)],
)
fastvideo/tests/stages/test_qwen_image21_stages.py
+304
View File
@@ -0,0 +1,304 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen request and stage contracts with real batches and weight-free components.
Tensor operations run on CPU. These checks exercise stage integration; independent
checkpoint parity remains in tests/local_tests/qwen_image21/.
"""
from types import SimpleNamespace
import pytest
import torch
from PIL import Image
from torch import nn
from fastvideo.models.dits.qwen_image21 import QwenImage21KVCacheAllocationError
from fastvideo.pipelines.basic.qwen_image21 import stages
from fastvideo.pipelines.basic.qwen_image21.stages import STATE_KEY
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
@pytest.fixture
def stage_args(monkeypatch):
# Keep the real stages and caches; only replace runtime device/weight moves.
monkeypatch.setattr(stages, "get_local_torch_device", lambda: torch.device("cpu"))
monkeypatch.setattr(stages.pinned_offload, "load", lambda model, device, pin=False: model.to(device))
return SimpleNamespace(
tp_size=1,
sp_size=1,
num_gpus=1,
text_encoder_cpu_offload=False,
vae_cpu_offload=False,
dit_cpu_offload=False,
dit_layerwise_offload=False,
use_fsdp_inference=False,
pin_cpu_memory=False,
output_type="pil",
pipeline_config=SimpleNamespace(
dit_config=SimpleNamespace(arch_config=SimpleNamespace(in_channels=2)),
vae_tiling=False,
reference_resolution=32,
use_kv_cache=True,
kv_cache_device="cpu",
),
)
def _batch(**updates):
values = dict(
data_type="image",
prompt="draw",
height=64,
width=64,
num_frames=1,
num_inference_steps=2,
guidance_scale=1.0,
true_cfg_scale=1.0,
seed=42,
)
values.update(updates)
return ForwardBatch(**values)
class TinyProcessor:
tokenizer = SimpleNamespace(encode=lambda _text: [9])
def __init__(self):
self.images = []
def apply_chat_template(self, *_args, **_kwargs):
return [[1, 2]]
def __call__(self, text, images=None, **kwargs):
assert kwargs["padding_side"] == "left"
self.images = list(images or [])
last_id = 5 if "bad" in text[0] else 4
ids = torch.tensor([[1, 2, 3] + [9] * len(self.images) + [last_id]])
return dict(input_ids=ids, attention_mask=torch.ones_like(ids))
class TinyEncoder(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones(1))
def forward(self, input_ids, **_kwargs):
return input_ids.float().unsqueeze(-1).repeat(1, 1, 2)
class TinyVAE(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones(1))
self.config = SimpleNamespace(latents_mean=(0.0, 0.0), latents_std=(1.0, 1.0))
self.encoded_pixels = []
def disable_tiling(self):
pass
def encode(self, pixels):
self.encoded_pixels.append(pixels.clone())
latent = torch.nn.functional.avg_pool2d(pixels[:, [0, 3], 0], 16).unsqueeze(2)
return SimpleNamespace(latent_dist=SimpleNamespace(mode=lambda: latent))
def decode(self, latents, return_dict=False):
image = latents[:, :1].repeat_interleave(16, 3).repeat_interleave(16, 4)
return (torch.cat((image.repeat(1, 3, 1, 1, 1), torch.full_like(image, -0.5)), dim=1),)
class TinyScheduler:
config = dict(base_image_seq_len=256, max_image_seq_len=8192, base_shift=0.5, max_shift=0.9)
def set_timesteps(self, sigmas, device, mu):
self.timesteps = torch.tensor(sigmas, device=device) * 1000
def set_begin_index(self, index):
assert index == 0
def step(self, noise, timestep, latents, return_dict=False):
return (latents - noise * 0.1,)
class TinyTransformer(nn.Module):
"""Populate real prefix caches so ownership and cleanup checks inspect data."""
def __init__(self, fail_cache=False, model_error=None):
super().__init__()
self.config = SimpleNamespace(causal_condition=True)
self.transformer_blocks = nn.ModuleList([nn.Identity()])
self.calls = []
self.fail_cache = fail_cache
self.model_error = model_error
def forward(self, hidden_states, encoder_hidden_states, kv_cache_mode, kv_cache, **_kwargs):
self.calls.append(SimpleNamespace(mode=kv_cache_mode, cache=kv_cache))
condition = encoder_hidden_states[:, -1:, :]
if kv_cache_mode == "extract":
prefix = condition.unsqueeze(2)
kv_cache.get_layer(0).store(prefix, prefix)
kv_cache.prefix_len = 1
if self.fail_cache:
self.fail_cache = False
raise QwenImage21KVCacheAllocationError("simulated host allocation failure")
elif kv_cache_mode == "cached":
condition = kv_cache.get_layer(0).get()[1].squeeze(2)
if self.model_error is not None:
raise self.model_error
return (condition.expand_as(hidden_states),)
def _run(batch, args, transformer=None):
transformer = transformer if transformer is not None else TinyTransformer()
vae, processor, scheduler = TinyVAE(), TinyProcessor(), TinyScheduler()
pipeline = [
stages.QwenImage21InputStage(),
stages.QwenImage21TextEncodingStage(TinyEncoder(), processor),
stages.QwenImage21LatentStage(vae),
stages.QwenImage21ScheduleStage(scheduler),
stages.QwenImage21DenoisingStage(transformer, scheduler),
stages.QwenImage21DecodingStage(vae),
]
for stage in pipeline:
batch = stage.forward(batch, args)
return batch, transformer, vae, processor
def _assert_caches_cleared(transformer):
caches = [call.cache for call in transformer.calls if call.cache is not None]
assert caches
for cache in caches:
assert cache.prefix_len is None
assert all(layer.k is None and layer.v is None for layer in cache.layer_caches)
def test_request_validation_accepts_real_forward_batch(stage_args):
batch = _batch()
result = stages.QwenImage21InputStage().forward(batch, stage_args)
assert result is batch
assert result.batch_size == 1
assert result.extra[STATE_KEY].img_shapes == [[(1, 4, 4)]]
assert not result.do_classifier_free_guidance
for updates, message in [
({"num_frames": 2}, "num_frames"),
({"guidance_scale": 3}, "true_cfg_scale"),
({"height": 63}, "multiple of 32"),
({"true_cfg_scale": float("nan")}, "finite"),
({"num_videos_per_prompt": 2}, "num_videos_per_prompt"),
]:
with pytest.raises(ValueError, match=message):
stages.QwenImage21InputStage().forward(_batch(**updates), stage_args)
for resolution in (0, 33):
stage_args.pipeline_config.reference_resolution = resolution
with pytest.raises(ValueError, match="reference_resolution"):
stages.QwenImage21InputStage().forward(_batch(), stage_args)
def test_ordered_references_preserve_rgba_for_vae_and_composite_for_encoder(stage_args):
references = [Image.new("RGBA", (32, 32), color) for color in [(0, 0, 0, 128), (128, 0, 0, 64)]]
batch, transformer, vae, processor = _run(_batch(references=references), stage_args)
assert batch.extra[STATE_KEY].img_shapes == [[(1, 2, 2), (1, 2, 2), (1, 4, 4)]]
assert batch.image_latent.shape == (1, 8, 2)
expected_red = torch.tensor([0.0, 128.0]) / 127.5 - 1
expected_alpha = torch.tensor([128.0, 64.0]) / 127.5 - 1
torch.testing.assert_close(batch.image_latent[0, ::4, 0], expected_red)
torch.testing.assert_close(batch.image_latent[0, ::4, 1], expected_alpha)
assert all(pixels.shape == (1, 4, 1, 32, 32) for pixels in vae.encoded_pixels)
for source, encoded in zip(references, processor.images, strict=True):
expected = Image.new("RGB", source.size, "white")
expected.paste(source, mask=source.getchannel("A"))
assert encoded.mode == "RGB"
assert encoded.tobytes() == expected.tobytes()
assert batch.output.shape == (1, 4, 1, 64, 64)
torch.testing.assert_close(batch.output[:, 3], torch.full((1, 1, 64, 64), 0.25))
_assert_caches_cleared(transformer)
def test_true_cfg_requires_explicit_negative_prompt_and_separate_caches(stage_args):
initial = torch.zeros(1, 16, 2)
guided, transformer, _, _ = _run(
_batch(latents=initial.clone(), true_cfg_scale=2, negative_prompt="bad"), stage_args)
assert guided.do_classifier_free_guidance
torch.testing.assert_close(guided.latents, torch.full_like(initial, -0.6))
assert [call.mode for call in transformer.calls] == ["extract", "extract", "cached", "cached"]
positive_cache, negative_cache = transformer.calls[0].cache, transformer.calls[1].cache
assert positive_cache is not negative_cache
assert transformer.calls[2].cache is positive_cache
assert transformer.calls[3].cache is negative_cache
_assert_caches_cleared(transformer)
unguided, transformer, _, _ = _run(_batch(latents=initial.clone(), true_cfg_scale=2), stage_args)
assert not unguided.do_classifier_free_guidance
assert [call.mode for call in transformer.calls] == ["extract", "cached"]
torch.testing.assert_close(unguided.latents, torch.full_like(initial, -0.8))
_assert_caches_cleared(transformer)
def test_host_cache_allocation_failure_falls_back_to_uncached_execution(stage_args):
transformer = TinyTransformer(fail_cache=True)
actual, _, _, _ = _run(_batch(), stage_args, transformer)
assert [call.mode for call in transformer.calls] == ["extract", None, None]
assert all(call.cache is None for call in transformer.calls[1:])
_assert_caches_cleared(transformer)
stage_args.pipeline_config.use_kv_cache = False
expected, _, _, _ = _run(_batch(), stage_args)
torch.testing.assert_close(actual.latents, expected.latents, atol=0, rtol=0)
torch.testing.assert_close(actual.output, expected.output, atol=0, rtol=0)
def test_model_error_propagates_and_clears_request_cache(stage_args):
error = RuntimeError("model failure")
transformer = TinyTransformer(model_error=error)
with pytest.raises(RuntimeError, match="model failure") as raised:
_run(_batch(), stage_args, transformer)
assert raised.value is error
assert [call.mode for call in transformer.calls] == ["extract"]
_assert_caches_cleared(transformer)
@pytest.mark.parametrize("offload", ["whole", "layerwise"])
def test_cancellation_stops_remaining_steps_and_releases_cache_and_weights(stage_args, monkeypatch, offload):
batch = _batch(num_inference_steps=3, latents=torch.zeros(1, 16, 2))
scheduler = TinyScheduler()
for stage in (
stages.QwenImage21InputStage(),
stages.QwenImage21TextEncodingStage(TinyEncoder(), TinyProcessor()),
stages.QwenImage21LatentStage(TinyVAE()),
stages.QwenImage21ScheduleStage(scheduler),
):
batch = stage.forward(batch, stage_args)
class InterruptingTransformer(TinyTransformer):
def forward(self, **kwargs):
result = super().forward(**kwargs)
denoising.interrupt = True
return result
transformer = InterruptingTransformer()
denoising = stages.QwenImage21DenoisingStage(transformer, scheduler)
stage_args.dit_cpu_offload = offload == "whole"
stage_args.dit_layerwise_offload = offload == "layerwise"
loaded, unloaded, layerwise_released, steps = [], [], [], []
monkeypatch.setattr(stages.pinned_offload, "load", lambda model, device, pin=False: loaded.append(model))
monkeypatch.setattr(stages.pinned_offload, "unload", lambda model: unloaded.append(model))
monkeypatch.setattr(denoising, "_release_layerwise", lambda: layerwise_released.append(True))
original_step = scheduler.step
def record_step(*args, **kwargs):
steps.append(args[1])
return original_step(*args, **kwargs)
monkeypatch.setattr(scheduler, "step", record_step)
result = denoising.forward(batch, stage_args)
assert result is batch
assert denoising.interrupt is True # The executor owns the cancellation latch.
assert [call.mode for call in transformer.calls] == ["extract"]
assert len(steps) == 1
assert batch.step_index == 0
torch.testing.assert_close(batch.latents, torch.full_like(batch.latents, -0.4))
_assert_caches_cleared(transformer)
assert loaded == unloaded == ([transformer] if offload == "whole" else [])
assert layerwise_released == ([True] if offload == "layerwise" else [])
fastvideo/tests/transformers/test_qwen_image21_attention.py
+410
View File
@@ -0,0 +1,410 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 attention and prefix-cache regressions with tiny native weights.
RoPE checks run on CPU; attention/cache cases require one CUDA GPU and the
ordinary FastVideo runtime. No model download or optional FlashAttention
installation is needed. Independent checkpoint parity remains in
``tests/local_tests/transformers/``.
"""
from __future__ import annotations
from functools import cache
import pytest
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.testing import assert_close
import fastvideo.envs as envs
import fastvideo.models.dits.qwen_image21 as qwen
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.platforms import AttentionBackendEnum
@pytest.fixture
def qwen_setup(env_overrides, monkeypatch, request):
if not torch.cuda.is_available():
pytest.skip("Requires one CUDA GPU")
env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False)
request.getfixturevalue("distributed_setup")
with torch.no_grad(), sdpa_kernel(SDPBackend.MATH):
yield torch.device("cuda")
def _initialize_linears(module):
# ReplicatedLinear allocates empty storage instead of initializing weights.
for layer in module.modules():
if isinstance(layer, ReplicatedLinear):
torch.nn.init.xavier_uniform_(layer.weight)
if layer.bias is not None:
torch.nn.init.zeros_(layer.bias)
return module.eval()
def _model(device):
config = QwenImage21Config(arch_config=QwenImage21ArchConfig(
in_channels=4, out_channels=4, num_layers=2, num_attention_heads=2,
attention_head_dim=8, context_in_dim=12, axes_dims_rope=(2, 2, 4), mlp_ratio=2))
return _initialize_linears(qwen.QwenImage21Transformer2DModel(config).to(device=device, dtype=torch.float32))
def _inputs(device):
# Adjacent reference slots must still form distinct bidirectional blocks.
return dict(
hidden_states=torch.randn(2, 12, 4, device=device),
encoder_hidden_states=torch.randn(2, 5, 12, device=device),
timestep=torch.tensor([0.6, 0.7], device=device),
img_shapes=[[(1, 2, 2)] * 3 for _ in range(2)],
img_mask=torch.tensor([[False, True, True, False, False, True]] * 2, device=device),
encoder_hidden_states_mask=torch.tensor(
[[True, True, True, True, False], [False, True, True, True, False]], device=device),
)
@pytest.mark.parametrize("position", [8191, 8192, 9215, 9216])
def test_rope_positive_position_does_not_alias_negative_coordinates(position):
rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 2))
mask = torch.tensor([True] + [False] * position)
actual = rope([(1, 1, 1)], mask, torch.device("cpu"))
# Each two-dimensional axis has frequency 1: its phase is exp(i * position).
expected = torch.exp(1j * torch.tensor([position] * 3, dtype=torch.float32))
assert_close(actual[-1], expected, atol=1e-6, rtol=1e-6)
image_expected = torch.exp(1j * torch.tensor([0, -1, -1], dtype=torch.float32))
assert_close(actual[0], image_expected, atol=1e-6, rtol=1e-6)
assert actual.shape == (position + 1, 3)
assert actual.dtype == torch.complex64 and actual.device.type == "cpu"
@pytest.mark.parametrize("device_name", ["cpu", "cuda"])
def test_rope_long_layout_preserves_multi_image_coordinates_and_short_requests(device_name):
if device_name == "cuda" and not torch.cuda.is_available():
pytest.skip("CUDA RoPE placement requires one GPU")
device = torch.device(device_name)
rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 2)).to(device=device, dtype=torch.bfloat16)
shapes = [(1, 2049, 1), (1, 1, 2049)]
def layout(text_tokens):
mask = torch.tensor([False] * text_tokens + [True] * 2049 + [False] * 3 + [True] * 2049 + [False] * 2,
device=device)
# Odd image axes span -1025..1023, beyond the old negative table.
# Each image advances subsequent text by its largest spatial axis.
positions = ([[p] * 3 for p in range(text_tokens)]
+ [[text_tokens, h, -1] for h in range(-1025, 1024)]
+ [[text_tokens + p] * 3 for p in (2049, 2050, 2051)]
+ [[text_tokens + 2052, -1, w] for w in range(-1025, 1024)]
+ [[text_tokens + p] * 3 for p in (4101, 4102)])
expected = torch.exp(1j * torch.tensor(positions, dtype=torch.float32, device=device))
return mask, expected
short_mask, short_expected = layout(3)
before = rope(shapes, short_mask, device).clone()
assert_close(before, short_expected, atol=1e-6, rtol=1e-6)
long_mask, long_expected = layout(9216)
actual = rope(shapes, long_mask, device)
assert_close(actual, long_expected, atol=1e-6, rtol=1e-6)
assert actual.device == long_mask.device and actual.dtype == torch.complex64
assert_close(rope(shapes, short_mask, device), before, atol=0, rtol=0)
@pytest.mark.parametrize("device_name", ["cpu", "cuda"])
def test_rope_long_layout_keeps_nontrivial_short_frequencies_bitwise(device_name):
if device_name == "cuda" and not torch.cuda.is_available():
pytest.skip("CUDA RoPE placement requires one GPU")
device = torch.device(device_name)
rope = qwen.QwenImage21Rope(theta=10000, axes_dim=(2, 2, 4)).to(device=device, dtype=torch.bfloat16)
shapes = [(1, 2, 2)]
short_mask = torch.tensor([False] * 3 + [True] * 4, device=device)
before = rope(shapes, short_mask, device).clone()
rope(shapes, torch.tensor([False] * 9216 + [True] * 4, device=device), device)
after = rope(shapes, short_mask, device)
assert after.dtype == torch.complex64 and after.device == short_mask.device
assert_close(after, before, atol=0, rtol=0)
def test_segmented_attention_matches_independent_dense_mask(qwen_setup):
device = qwen_setup
attention = _initialize_linears(qwen.QwenImage21Attention(16, 2, 8, 1e-6).to(device))
hidden = torch.randn(2, 15, 16, device=device)
image_ids = torch.tensor([-1, -1] + [0] * 4 + [1] * 4 + [-1] + [2] * 4, device=device)
valid = torch.ones(2, 15, dtype=torch.bool, device=device)
valid[0, 0] = False
valid[1, 1] = False
# Independently encode causality: preceding tokens, or any token in the
# same image. In particular, adjacent reference images cannot see ahead.
positions = torch.arange(15, device=device)
causal = positions[:, None] >= positions[None, :]
same_image = (image_ids[:, None] == image_ids[None, :]) & (image_ids[:, None] >= 0)
mask = (causal | same_image)[None, None] & valid[:, None, None, :]
query = attention.norm_q(attention.to_q(hidden)[0].unflatten(-1, (2, 8)))
key = attention.norm_k(attention.to_k(hidden)[0].unflatten(-1, (2, 8)))
value = attention.to_v(hidden)[0].unflatten(-1, (2, 8))
expected = F.scaled_dot_product_attention(
query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=mask, dropout_p=0.0)
expected = expected.masked_fill(~mask.any(-1, keepdim=True), 0).transpose(1, 2).flatten(2, 3)
expected = attention.to_out[0](expected)[0]
actual = attention(hidden, segments=qwen._qwenimage21_prefix_segments(image_ids, 11), key_valid=valid)
assert torch.isfinite(actual).all()
assert_close(actual, expected, atol=1e-5, rtol=1e-5)
def test_cached_decode_matches_recompute_and_owns_prefix_storage(qwen_setup):
device = qwen_setup
model = _model(device)
inputs = _inputs(device)
cache = qwen.QwenImage21KVCache(model.config.num_layers, storage_device="cpu")
extracted = model(**inputs, kv_cache=cache, kv_cache_mode="extract")
assert_close(extracted, model(**inputs), atol=0, rtol=0)
assert cache.prefix_len == 11
prefix_values = []
for layer in cache.layer_caches:
assert layer.k.shape == layer.v.shape == (2, 11, 2, 8)
assert layer.k.device.type == layer.v.device.type == "cpu"
for tensor in (layer.k, layer.v):
assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size()
prefix_values.append((layer.k.clone(), layer.v.clone()))
inputs["timestep"] = torch.tensor([0.25, 0.35], device=device)
inputs["hidden_states"] = inputs["hidden_states"].clone()
inputs["hidden_states"][:, -4:] += 0.7
actual = model(**inputs, kv_cache=cache, kv_cache_mode="cached")
expected = model(**inputs)
assert actual.shape == (2, 4, 4)
assert_close(actual, expected[:, -4:], atol=1e-5, rtol=1e-5)
assert_close(extracted[:, :-4], expected[:, :-4], atol=0, rtol=0)
for layer, (key, value) in zip(cache.layer_caches, prefix_values, strict=True):
assert_close(layer.k, key, atol=0, rtol=0)
assert_close(layer.v, value, atol=0, rtol=0)
# A prefix view must own its bytes, even when source and cache are both on
# the CPU. Otherwise retaining it keeps the entire prefill allocation alive.
source = torch.randn(1, 15, 2, 8)
owned = qwen.QwenImage21KVLayerCache("cpu")
owned.store(source[:, :3], source[:, :3])
original = source[:, :3].clone()
assert owned.k.untyped_storage().data_ptr() != source.untyped_storage().data_ptr()
assert owned.k.untyped_storage().nbytes() == owned.k.numel() * owned.k.element_size()
source.zero_()
assert_close(owned.k, original, atol=0, rtol=0)
assert_close(owned.v, original, atol=0, rtol=0)
cache.clear()
owned.clear()
assert cache.prefix_len is None
assert all(layer.k is None and layer.v is None for layer in cache.layer_caches)
with pytest.raises(RuntimeError, match="extract"):
owned.get()
def test_padding_embeddings_cannot_change_target_predictions(qwen_setup):
device = qwen_setup
model = _model(device)
inputs = _inputs(device)
expected = model(**inputs)[:, -4:]
embeddings = inputs["encoder_hidden_states"].clone()
embeddings[~inputs["encoder_hidden_states_mask"]] = 1000
inputs["encoder_hidden_states"] = embeddings
cache = qwen.QwenImage21KVCache(model.config.num_layers, storage_device="cpu")
extracted = model(**inputs, kv_cache=cache, kv_cache_mode="extract")[:, -4:]
cached = model(**inputs, kv_cache=cache, kv_cache_mode="cached")
assert torch.isfinite(extracted).all()
assert torch.isfinite(cached).all()
assert_close(extracted, expected, atol=0, rtol=0)
assert_close(cached, expected, atol=1e-5, rtol=1e-5)
cache.clear()
def test_masked_calls_use_sdpa_and_dense_calls_use_selected_backend(qwen_setup, monkeypatch):
import fastvideo.attention.selector as selector
device = qwen_setup
dense_calls, resolutions, masked_calls = [], [], []
class RecordingImpl:
def __init__(self, **kwargs):
assert kwargs["causal"] is False
assert kwargs["num_heads"] == kwargs["num_kv_heads"] == 2
assert kwargs["head_size"] == 8
def forward(self, query, key, value, attn_metadata):
assert attn_metadata is None
dense_calls.append((query.shape[1], key.shape[1]))
return F.scaled_dot_product_attention(
query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2),
dropout_p=0.0).transpose(1, 2)
class RecordingBackend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def get_impl_cls():
return RecordingImpl
def resolve_backend(head_size, dtype, **kwargs):
assert head_size == 8
resolutions.append(kwargs)
return RecordingBackend
monkeypatch.setattr(selector, "get_attn_backend", resolve_backend)
baseline = _initialize_linears(qwen.QwenImage21Attention(
16, 2, 8, 1e-6, requested_backend=AttentionBackendEnum.TORCH_SDPA).to(device))
assert resolutions == []
attention = qwen.QwenImage21Attention(16, 2, 8, 1e-6).to(device).eval()
explicit_attention = qwen.QwenImage21Attention(
16, 2, 8, 1e-6, requested_backend=AttentionBackendEnum.FLASH_ATTN).to(device).eval()
attention.load_state_dict(baseline.state_dict())
explicit_attention.load_state_dict(baseline.state_dict())
assert [resolution.get("requested", selector.NO_REQUEST) for resolution in resolutions] == [
selector.NO_REQUEST, AttentionBackendEnum.FLASH_ATTN]
assert all("default_backend" not in resolution for resolution in resolutions)
assert all(AttentionBackendEnum.FLASH_ATTN in resolution["supported_attention_backends"]
for resolution in resolutions)
hidden = torch.randn(1, 11, 16, device=device)
image_ids = torch.tensor([-1, -1] + [0] * 4 + [-1] + [1] * 4, device=device)
segments = qwen._qwenimage21_prefix_segments(image_ids, 7)
valid = torch.ones(1, 11, dtype=torch.bool, device=device)
valid[:, 0] = False
expected = baseline(hidden, segments=segments)
padded_expected = baseline(hidden, segments=segments, key_valid=valid)
assert_close(explicit_attention(hidden, segments=segments), expected, atol=1e-5, rtol=1e-5)
assert dense_calls == [(4, 6), (4, 11)]
dense_calls.clear()
sdpa = qwen._attention
def record_masked_sdpa(query, key, value, mask=None):
assert mask is not None
masked_calls.append((query.shape[1], key.shape[1]))
return sdpa(query, key, value, mask)
monkeypatch.setattr(qwen, "_attention", record_masked_sdpa)
cache = qwen.QwenImage21KVLayerCache()
actual = attention(hidden, segments=segments, layer_cache=cache,
kv_cache_mode="extract", cache_write_slice=slice(0, 7))
assert_close(actual, expected, atol=1e-5, rtol=1e-5)
assert masked_calls == [(2, 2), (1, 7)]
assert dense_calls == [(4, 6), (4, 11)]
dense_calls.clear()
masked_calls.clear()
cached = attention(hidden[:, 7:], layer_cache=cache, kv_cache_mode="cached")
assert_close(cached, expected[:, 7:], atol=1e-5, rtol=1e-5)
assert dense_calls == [(4, 11)]
assert masked_calls == []
dense_calls.clear()
padded = attention(hidden, segments=segments, key_valid=valid)
assert_close(padded, padded_expected, atol=1e-5, rtol=1e-5)
assert dense_calls == []
assert masked_calls == [(2, 2), (4, 6), (1, 7), (4, 11)]
masked_calls.clear()
padded_cached = attention(hidden[:, 7:], layer_cache=cache, kv_cache_mode="cached",
attention_mask=valid[:, None, None, :])
assert_close(padded_cached, padded_expected[:, 7:], atol=1e-5, rtol=1e-5)
assert dense_calls == []
assert masked_calls == [(4, 11)]
cache.clear()
@pytest.fixture
def selector_resolution(qwen_setup, monkeypatch):
import fastvideo.attention.selector as selector
import fastvideo.platforms as platforms
from fastvideo.attention.backends.sdpa import SDPABackend
class AlternativeImpl:
def __init__(self, **kwargs):
pass
class AlternativeBackend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def get_impl_cls():
return AlternativeImpl
def platform_backend(requested, head_size, dtype):
# A dense alternative is available automatically, so losing an SDPA
# override cannot silently pass by falling back to SDPA anyway.
return "TORCH_SDPA" if requested is AttentionBackendEnum.TORCH_SDPA else "FLASH_ATTN"
monkeypatch.setattr(platforms.current_platform, "get_attn_backend_cls", platform_backend)
monkeypatch.setattr(selector, "resolve_obj_by_qualname",
lambda name: SDPABackend if name == "TORCH_SDPA" else AlternativeBackend)
# Exercise the real selector and its cache without leaving fake platform
# resolutions in the shared process cache or clearing unrelated entries.
monkeypatch.setattr(selector, "_cached_get_attn_backend", cache(selector._cached_get_attn_backend.__wrapped__))
def _direct_attention(kind, recorded_backend=None):
if kind == "block":
return qwen.QwenImage21TransformerBlock(16, 2, 8, 2, 1e-6).attn
config = QwenImage21Config(arch_config=QwenImage21ArchConfig(
in_channels=4, out_channels=4, num_layers=1, num_attention_heads=2,
attention_head_dim=8, context_in_dim=12, axes_dims_rope=(2, 2, 4), mlp_ratio=2))
config._resolved_attention_backend = recorded_backend
return qwen.QwenImage21Transformer2DModel(config).transformer_blocks[0].attn
def _selected_backend(attention):
return (AttentionBackendEnum.TORCH_SDPA if attention.dense_attn is None else attention.dense_attn.backend)
@pytest.mark.parametrize("kind", ["transformer", "block"])
@pytest.mark.parametrize("environment", ["TORCH_SDPA", "FLASH_ATTN", None])
def test_direct_construction_honors_environment_backend(selector_resolution, kind, environment):
with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
attention = _direct_attention(kind)
expected = AttentionBackendEnum.TORCH_SDPA if environment == "TORCH_SDPA" else AttentionBackendEnum.FLASH_ATTN
assert _selected_backend(attention) is expected
@pytest.mark.parametrize("recorded", [AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN])
def test_direct_transformer_recorded_backend_overrides_environment(selector_resolution, recorded):
from fastvideo.attention.selector import _component_attention_backend_scope
environment = "FLASH_ATTN" if recorded is AttentionBackendEnum.TORCH_SDPA else "TORCH_SDPA"
with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
with _component_attention_backend_scope(environment, component="qwen"):
attention = _direct_attention("transformer", recorded)
assert _selected_backend(attention) is recorded
@pytest.mark.parametrize("requested", [AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FLASH_ATTN])
def test_direct_block_explicit_backend_overrides_environment(selector_resolution, requested):
environment = "FLASH_ATTN" if requested is AttentionBackendEnum.TORCH_SDPA else "TORCH_SDPA"
with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
block = qwen.QwenImage21TransformerBlock(16, 2, 8, 2, 1e-6, requested_backend=requested)
assert _selected_backend(block.attn) is requested
@pytest.mark.parametrize("kind", ["transformer", "block"])
@pytest.mark.parametrize("scoped,consult_env,environment,expected", [
(AttentionBackendEnum.TORCH_SDPA, False, "FLASH_ATTN", AttentionBackendEnum.TORCH_SDPA),
(None, False, "TORCH_SDPA", AttentionBackendEnum.FLASH_ATTN),
(None, True, "TORCH_SDPA", AttentionBackendEnum.TORCH_SDPA),
])
def test_direct_construction_honors_component_scope(selector_resolution, kind, scoped, consult_env, environment, expected):
from fastvideo.attention.selector import _component_attention_backend_scope
with envs.FASTVIDEO_ATTENTION_BACKEND.override(environment):
with _component_attention_backend_scope(scoped, component="qwen", consult_env=consult_env):
attention = _direct_attention(kind)
assert _selected_backend(attention) is expected
tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py
+173
View File
@@ -0,0 +1,173 @@
# SPDX-License-Identifier: Apache-2.0
"""Released-weight Qwen-Image-2.1 conditioner parity through the production loader.
Run with QWEN_IMAGE21_RUN_ENCODER_PARITY=1 and QWEN_IMAGE21_MODEL_ROOT set to
the complete local HF snapshot. This gate requires a BF16 CUDA GPU; opting in
turns missing assets or dependencies into failures rather than skipped parity.
Set QWEN_IMAGE21_ENCODER_PARITY_REPORT to save the numerical diagnostics as JSON.
"""
import gc
import json
import os
from pathlib import Path
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from torch.testing import assert_close
PARITY_SCOPE = "production_loader"
ATOL = 1e-2
RTOL = 1e-2
def _numerical_metrics(actual, expected):
actual, expected = actual.float(), expected.float()
difference = actual - expected
reference_rms = expected.square().mean().sqrt()
rmse = difference.square().mean().sqrt()
actual_flat, expected_flat = actual.flatten().double(), expected.flatten().double()
cosine = torch.nn.functional.cosine_similarity(actual_flat, expected_flat, dim=0).clamp(-1, 1)
return {
"shape": list(actual.shape),
"finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()),
"max_abs": difference.abs().max().item(),
"mean_abs": difference.abs().mean().item(),
"rmse": rmse.item(),
"reference_rms": reference_rms.item(),
"relative_l2": (rmse / reference_rms.clamp_min(torch.finfo(torch.float32).tiny)).item(),
"cosine": cosine.item(),
"mismatched_elements": int((difference.abs() > ATOL + RTOL * expected.abs()).sum()),
"total_elements": actual.numel(),
}
def _make_cases(processor):
prompt = "<|im_start|>user\nA red fox on fresh snow.<|im_end|>\n<|im_start|>assistant\n"
cases = {"text": processor(text=[prompt], padding=True, padding_side="left", return_tensors="pt")}
long_prompt = "<|im_start|>user\n" + "Small red fox on fresh snow. " * 220 + "<|im_end|>\n"
cases["long_text"] = processor(text=[long_prompt], return_tensors="pt")
assert cases["long_text"]["input_ids"].shape[-1] > 1024
for count in (1, 2, 10):
images = []
for index in range(count):
pixels = np.arange(64 * 96 * 3, dtype=np.uint32).reshape(64, 96, 3)
images.append(Image.fromarray(((pixels * (index + 1) + index * 31) % 256).astype(np.uint8)))
image_prompt = "<|im_start|>user\n" + " ".join(
f"<image{index + 1}><|vision_start|><|image_pad|><|vision_end|>" for index in range(count))
image_prompt += " Blend these references into a snowy scene.<|im_end|>\n"
cases[f"images_{count}"] = processor(
text=[image_prompt],
images=images,
min_pixels=4096,
max_pixels=16384,
return_tensors="pt",
)
cases["left_padding"] = processor(text=[prompt, prompt + " A blue sky."], padding=True,
padding_side="left", return_tensors="pt")
return cases
def test_qwen_image21_conditioner_real_weight_parity():
if os.environ.get("QWEN_IMAGE21_RUN_ENCODER_PARITY") != "1":
pytest.skip("set QWEN_IMAGE21_RUN_ENCODER_PARITY=1 on an allocated CUDA GPU")
if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
pytest.fail("Qwen-Image-2.1 conditioner parity requires a BF16 CUDA GPU", pytrace=False)
checkpoint = os.environ.get("QWEN_IMAGE21_MODEL_ROOT")
if not checkpoint:
pytest.fail("set QWEN_IMAGE21_MODEL_ROOT to the downloaded Qwen-Image-2.1 HF snapshot", pytrace=False)
root = Path(checkpoint)
if not (root / "text_encoder" / "config.json").is_file() or not (root / "processor").is_dir():
pytest.fail("The checkpoint must include text_encoder and processor component directories", pytrace=False)
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor, __version__ as transformers_version
from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLConfig
from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
from fastvideo.models.loader.component_loader import TextEncoderLoader
device = torch.device("cuda", 0)
processor = Qwen3VLProcessor.from_pretrained(root / "processor", local_files_only=True)
cases = _make_cases(processor)
official, loading_info = Qwen3VLForConditionalGeneration.from_pretrained(
root / "text_encoder",
dtype=torch.bfloat16,
attn_implementation="sdpa",
local_files_only=True,
low_cpu_mem_usage=True,
output_loading_info=True,
)
errors = {name: loading_info[name] for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
if loading_info.get(name)}
assert not errors, f"Official checkpoint strict loading failed: {errors}"
official = official.to(device).eval()
norm = official.model.language_model.norm
handle = norm.register_forward_hook(lambda _module, args, _output: args[0])
expected = {}
try:
for name, case in cases.items():
inputs = {key: value.to(device=device, dtype=torch.bfloat16 if key == "pixel_values" else value.dtype)
for key, value in case.items()}
with torch.no_grad():
output = official(**inputs, output_hidden_states=True, use_cache=False)
expected[name] = output.hidden_states[-1].detach().cpu()
del output, inputs
finally:
handle.remove()
del official, norm
gc.collect()
torch.cuda.empty_cache()
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29625")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")
maybe_init_distributed_environment_and_model_parallel(1, 1)
config = QwenImage21Qwen3VLConfig()
config.update_model_arch(json.loads((root / "text_encoder" / "config.json").read_text()))
args = SimpleNamespace(
text_encoder_cpu_offload=True,
override_text_encoder_quant=None,
override_text_encoder_safetensors=None,
pin_cpu_memory=False,
disable_offload_on_unified_memory=lambda _device_id, offload_flag=None: False,
)
report = {
"component": "qwen_image21_conditioner",
"checkpoint_revision": root.name,
"torch_version": torch.__version__,
"transformers_version": transformers_version,
"atol": ATOL,
"rtol": RTOL,
"cases": {},
}
failures = []
try:
native = TextEncoderLoader().load_model(str(root / "text_encoder"), config, device, args, dtype="bf16")
for name, case in cases.items():
inputs = {key: value.to(device=device, dtype=torch.bfloat16 if key == "pixel_values" else value.dtype)
for key, value in case.items()}
with torch.no_grad():
actual = native(**inputs).detach().cpu()
mask = case["attention_mask"].bool()
metrics = _numerical_metrics(actual[mask], expected[name][mask])
report["cases"][name] = metrics
print(f"QWEN_IMAGE21_ENCODER_PARITY {name} {json.dumps(metrics, sort_keys=True)}", flush=True)
try:
assert_close(actual[mask], expected[name][mask], rtol=RTOL, atol=ATOL,
msg=lambda message: f"{name}: {message}")
except AssertionError as error:
failures.append(f"{name}: {error}")
del actual, inputs
finally:
report_path = os.environ.get("QWEN_IMAGE21_ENCODER_PARITY_REPORT")
if report_path:
destination = Path(report_path)
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
cleanup_dist_env_and_memory()
assert not failures, "\n\n".join(failures)
tests/local_tests/qwen_image21/README.md
+138
View File
@@ -0,0 +1,138 @@
# Qwen-Image-2.1 validation
Use the package tests for regression coverage and the opt-in local tests to
compare FastVideo with pinned upstream Diffusers. For generation commands and
memory measurements, see the [cookbook](../../../docs/cookbook/qwen-image21.md).
## Confirmed coverage
Validation on NVIDIA GB200 at `c3d59459` completed on 2026-10-23: **36 cases
passed, with no failures or skips**.
| Check | Cases | Result |
| --- | ---: | --- |
| Package attention and stage tests | 9 | Backend selection, masks, caches and pipeline stages pass |
| FA4 transformer-block golden | 1 | Zero drift against the reviewed device/runtime reference |
| FA4 T2I and editing image regressions | 2 | SSIM 1.0; unchanged threshold 0.98 |
| Encoder, DiT and VAE component parity | 12 | All component cases pass |
| SDPA upstream pipeline parity | 12 | All modes below pass at 40 steps |
The pipeline matrix covers T2I, plain I2I, editing, reference-to-image,
transparent RGBA generation and editing, annotations, mask images, true CFG,
uncached execution and ten references at 1024 x 1024. A separate T2I case covers
2048 x 2048. Conditioning, initial/final latents, all 40 noise predictions and
latent updates, floating-point RGBA output and quantized PNG pixels were exact.
Raw FP32 scheduler arrays have maximum differences of `1.1920928955078125e-7`
for sigmas and `0.0001220703125` for timesteps; they are not bitwise equal.
Upstream comparisons explicitly use SDPA with matching cuDNN SDPA settings.
FA4 is validated separately against its golden and reviewed image references;
FA4 and SDPA outputs are not expected to be bitwise equal. The automatic
attention selector was also observed selecting FA4 for all 64 blocks across the
two image regressions with `FASTVIDEO_FA4=1`.
These results establish the tested GB200 configurations. They do not establish
other GPU/backend combinations, 2K with ten references, or hard preservation of
pixels outside a semantic mask.
## Package regression tests
Install FastVideo using the [CUDA installation guide](../../../docs/getting_started/installation/gpu.md)
and run commands from the repository root on an allocated CUDA GPU:
```bash
python -m pytest fastvideo/tests/transformers/test_qwen_image21_attention.py \
fastvideo/tests/stages/test_qwen_image21_stages.py -xq
FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python -m pytest \
fastvideo/tests/golden_gate/test_qwen_image21.py -xq
FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python -m pytest \
fastvideo/tests/ssim/test_qwen_image21_similarity.py -xq
```
These files cover attention, pipeline stages, transformer blocks and generated
images through the existing package CI lanes. The golden loads one released
BF16 transformer block. The image tests generate T2I and editing outputs at
512 x 512, four steps and seed 21, comparing all four RGBA channels.
The shared harnesses download missing references from
[`FastVideo/ssim-reference-videos`](https://huggingface.co/datasets/FastVideo/ssim-reference-videos).
Qwen references are available for the GB200 FA4 golden and default image tier.
The optional `--ssim-full-quality` tier uses 1024 x 1024 and 40 steps and needs
its own reviewed references; those are not covered by the published default
tier. See the [SSIM guide](../../../fastvideo/tests/ssim/README.md) for reference
management.
## Prepare upstream parity
Local parity tests under `tests/local_tests/` are opt-in and run outside CI.
Use BF16 on a CUDA GPU, a complete Qwen checkpoint and a clean checkout of the
pinned upstream revision. The pipeline harness uses CPU offload, VAE tiling
and CPU prefix KV storage; its pytest duration is not a generation benchmark.
Use the cookbook's measured memory profiles to plan resources.
```bash
git clone https://github.com/huggingface/diffusers.git ../diffusers-qwen-image21
git -C ../diffusers-qwen-image21 checkout 578c9b2c6636ab2424a0e56186268b83623656b2
hf download Qwen/Qwen-Image-2.1 \
--revision d26bb61231c349cf6b7896fa83353113880e1ba3 \
--local-dir ../Qwen-Image-2.1
export QWEN_IMAGE21_DIFFUSERS_DIR="$PWD/../diffusers-qwen-image21"
export QWEN_IMAGE21_MODEL_ROOT="$PWD/../Qwen-Image-2.1"
export PYTHONPATH="$QWEN_IMAGE21_DIFFUSERS_DIR/src${PYTHONPATH:+:$PYTHONPATH}"
export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
```
An existing checkpoint can be used by changing `QWEN_IMAGE21_MODEL_ROOT`. It
must include `model_index.json`, `processor`, `text_encoder`, `transformer`,
`vae`, `scheduler` and all weight shards. No checkpoint conversion is needed.
Explicitly enabled real-weight tests fail if prerequisites are missing.
### Component parity
```bash
QWEN_IMAGE21_RUN_ENCODER_PARITY=1 python -m pytest \
tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py -xq
QWEN_IMAGE21_RUN_DIT_PARITY=1 python -m pytest \
tests/local_tests/transformers/test_qwen_image21_transformer_parity.py -xq
QWEN_IMAGE21_RUN_VAE_PARITY=1 python -m pytest \
tests/local_tests/vaes/test_qwen_image21_vae_parity.py -xq
```
These checks cover text and 1/2/10-image conditioning, every DiT block and
cached decoding, and VAE posterior/decode behavior including tiling. Optional
`QWEN_IMAGE21_ENCODER_PARITY_REPORT` and `QWEN_IMAGE21_DIT_PARITY_REPORT` paths
save component diagnostics as JSON.
### Pipeline parity
Run all eleven 1024-pixel modes and the separate 2K T2I case at 40 steps:
```bash
QWEN_IMAGE21_RUN_PIPELINE_PARITY=1 QWEN_IMAGE21_RUN_2K_PARITY=1 \
QWEN_IMAGE21_PARITY_STEPS=40 \
python -m pytest tests/local_tests/qwen_image21/test_pipeline_parity.py -xq \
--basetemp outputs/qwen_image21/parity-40step
```
Choose a fresh `--basetemp` path for each run: pytest replaces an existing
directory at that path. For one mode, select its node, for example
`tests/local_tests/qwen_image21/test_pipeline_parity.py::test_qwen_image21_modes_1024_pipeline_parity[edit]`.
Without `QWEN_IMAGE21_PARITY_STEPS`, the harness defaults to four steps.
Each case runs native and upstream pipelines sequentially in separate
processes, with identical BF16 initial latents, sigmas and cache settings.
Synthetic reference images are generated by default.
`QWEN_IMAGE21_PARITY_CASES_JSON` can supply per-case prompts and ordered image
paths; relative paths resolve beside that JSON file. For parallel invocations,
assign distinct `CUDA_VISIBLE_DEVICES`, `MASTER_PORT` and artifact paths.
The tests assert RGBA output parity within the existing BF16 tolerance and
verify PNG alpha export. They also write final-output error metrics. The stronger
exact-trajectory results reported above came from separate diagnostic runs;
a passing pixel assertion alone does not establish bitwise equality at every
stage.
tests/local_tests/qwen_image21/test_pipeline_parity.py
+235
View File
@@ -0,0 +1,235 @@
# SPDX-License-Identifier: Apache-2.0
"""Opt-in, released-weight pipeline parity through the public FastVideo API.
Each implementation runs in its own process so their weights never overlap in
VRAM. Synthetic references are the default; QWEN_IMAGE21_PARITY_CASES_JSON can
provide per-case prompts and ordered reference paths from the original source.
"""
from __future__ import annotations
import json
import os
import subprocess
import sys
from pathlib import Path
import numpy as np
import pytest
import torch
from PIL import Image, ImageDraw, ImageOps
ROOT = Path(__file__).resolve().parents[3]
REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
PARITY_SCOPE = "pipeline"
CASES = ("t2i", "i2i", "edit", "ref2img", "ten_refs", "rgba", "rgba_edit", "annotation", "mask_image", "true_cfg", "uncached")
def _worker(kind: str, spec_path: str, output_path: str) -> None:
# Native CUDA startup disables cuDNN SDPA. Match that policy in the isolated
# reference process; its default can change encoder numerics on GB200.
torch.backends.cuda.enable_cudnn_sdp(False)
spec = json.loads(Path(spec_path).read_text())
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(Path(spec["reference"]) / "src"))
height, width = spec.get("height", spec["resolution"]), spec.get("width", spec["resolution"])
seed = spec.get("seed", 21)
reference_resolution = spec.get("reference_resolution", 1024)
sigmas = spec.get("sigmas", np.linspace(1.0, 1 / spec["steps"], spec["steps"]).tolist())
noise = torch.load(spec["latents"], map_location="cpu", weights_only=True)
assert noise.dtype == torch.bfloat16
assert noise.shape == (1, (height // 16) * (width // 16), 64)
torch.manual_seed(seed)
if kind == "native":
from fastvideo import VideoGenerator
from fastvideo.api import (EngineConfig, GenerationRequest, GeneratorConfig, InputConfig, OffloadConfig,
OutputConfig, ParallelismConfig, PipelineSelection, SamplingConfig)
generator = VideoGenerator.from_config(
GeneratorConfig(
model_path=spec["model"],
engine=EngineConfig(num_gpus=1, parallelism=ParallelismConfig(tp_size=1, sp_size=1),
offload=OffloadConfig(dit=True, dit_layerwise=True, text_encoder=True,
vae=True, lazy_module_load=True)),
pipeline=PipelineSelection(workload_type="i2i" if spec["refs"] else "t2i",
experimental={"attention_backend": "TORCH_SDPA",
"reference_resolution": reference_resolution,
"use_kv_cache": spec["cache"]}),
))
try:
result = generator.generate(
GenerationRequest(
prompt=spec["prompt"], negative_prompt=spec["negative_prompt"],
inputs=InputConfig(references=spec["refs"] or None, latents=noise.clone()),
sampling=SamplingConfig(height=height, width=width, num_frames=1,
fps=1, num_inference_steps=spec["steps"], guidance_scale=1,
true_cfg_scale=spec["true_cfg_scale"], seed=seed, sigmas=sigmas),
output=OutputConfig(output_path=str(Path(output_path).with_suffix(".png")),
return_frames=True, save_video=True),
))
torch.save({"pixels": result.samples.detach().float().cpu(), "png_path": result.video_path,
"peak_memory_mb": result.peak_memory_mb}, output_path)
finally:
generator.shutdown()
else:
from diffusers.pipelines.qwenimage21 import pipeline_qwenimage21
expected = Path(spec["reference"]) / "src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py"
assert Path(pipeline_qwenimage21.__file__).resolve() == expected.resolve()
pipe = pipeline_qwenimage21.QwenImage21Pipeline.from_pretrained(
spec["model"], torch_dtype=torch.bfloat16, local_files_only=True)
pipe.vae.enable_tiling()
pipe.enable_sequential_cpu_offload()
references = []
for path in spec["refs"]:
with Image.open(path) as image:
references.append(ImageOps.exif_transpose(image).convert("RGBA"))
with torch.no_grad():
output = pipe(prompt=spec["prompt"], image=references or None, height=height, width=width,
output_resolution=reference_resolution, num_inference_steps=spec["steps"], sigmas=sigmas,
negative_prompt=spec["negative_prompt"], true_cfg_scale=spec["true_cfg_scale"],
generator=torch.Generator(device="cpu").manual_seed(seed),
latents=noise.clone(), output_type="pt", use_kv_cache=spec["cache"])
torch.save({"pixels": output.images.unsqueeze(2).float().cpu()}, output_path)
def _assets(tmp_path: Path) -> list[str]:
files = []
for index in range(10):
image = Image.new("RGBA", (256, 256), (230, 235, 240, 0 if index == 0 else 255))
draw = ImageDraw.Draw(image)
draw.ellipse((48 + index, 48, 208, 208), fill=(30 + index * 20, 80, 140, 200))
if index == 1:
draw.rectangle((24, 24, 232, 232), outline="red", width=8)
path = tmp_path / f"reference_{index + 1}.png"
image.save(path)
files.append(str(path))
mask = Image.new("RGBA", (256, 256), "black")
ImageDraw.Draw(mask).ellipse((48, 48, 208, 208), fill="white")
path = tmp_path / "mask.png"
mask.save(path)
return files + [str(path)]
def _case_overrides(case: str) -> dict:
path = os.environ.get("QWEN_IMAGE21_PARITY_CASES_JSON")
if not path:
return {}
manifest = Path(path).resolve()
overrides = json.loads(manifest.read_text()).get(case, {})
allowed = {"prompt", "refs", "height", "width", "steps", "seed", "reference_resolution", "negative_prompt",
"true_cfg_scale", "cache", "sigmas"}
assert isinstance(overrides, dict) and not (set(overrides) - allowed)
if "refs" in overrides:
assert isinstance(overrides["refs"], list) and len(overrides["refs"]) <= 10
paths = [Path(value) for value in overrides["refs"]]
paths = [value if value.is_absolute() else manifest.parent / value for value in paths]
assert all(value.is_file() for value in paths), "Manifest references must be existing local files"
overrides["refs"] = [str(value.resolve()) for value in paths]
return overrides
def _difference_metrics(actual: torch.Tensor, expected: torch.Tensor, *, atol=0.04, rtol=0.02) -> dict[str, float]:
actual, expected = actual.float(), expected.float()
error = (actual - expected).abs()
mean_error = error.mean().item()
reference_mean = expected.abs().mean().item()
rmse = error.square().mean().sqrt().item()
return {
"abs_mean": mean_error,
"abs_max": error.max().item(),
"rmse": rmse,
"relative_l1": mean_error / max(reference_mean, 1e-12),
"relative_l2": rmse / max(expected.square().mean().sqrt().item(), 1e-12),
"actual_abs_mean": actual.abs().mean().item(),
"reference_abs_mean": reference_mean,
"fraction_outside_tolerance": (error > atol + rtol * expected.abs()).float().mean().item(),
}
def _run_parity(case, resolution, tmp_path):
if os.environ.get("QWEN_IMAGE21_RUN_PIPELINE_PARITY") != "1":
pytest.skip("Set QWEN_IMAGE21_RUN_PIPELINE_PARITY=1 on an allocated CUDA node")
if not torch.cuda.is_available():
pytest.fail("Qwen-Image-2.1 pipeline parity requires CUDA", pytrace=False)
model = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", ROOT / "official_weights/Qwen-Image-2.1"))
reference = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", ROOT / "official_reference/diffusers"))
if not (model / "model_index.json").is_file() or not (reference / "src").is_dir():
pytest.fail("Provide QWEN_IMAGE21_MODEL_ROOT and the pinned QWEN_IMAGE21_DIFFUSERS_DIR", pytrace=False)
revision = subprocess.check_output(["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip()
assert revision == REFERENCE_REVISION
subprocess.run(["git", "-C", str(reference), "diff", "--exit-code", "HEAD", "--", "src/diffusers"], check=True)
files = _assets(tmp_path)
counts = dict(i2i=1, edit=1, ref2img=2, ten_refs=10, rgba_edit=1, annotation=2, mask_image=1, true_cfg=1)
refs = files[:counts.get(case, 0)]
prompt = "An elegant blue ceramic vase in a studio, with the word Qwen written on it."
if refs:
prompt = "Make a blue sculpture inspired by the shapes in the reference images."
if case == "edit":
prompt = "Change the object in image 1 into a blue ceramic sphere; retain the background and composition."
if case in ("rgba", "rgba_edit"):
prompt = "A blue sphere isolated on a transparent background."
if case == "annotation":
prompt = "Change the object enclosed by the red annotation in image 2 into a blue sphere in image 1."
if case == "mask_image":
refs.append(files[-1])
prompt = "Replace the white masked area of image 2 in image 1 with flowers."
noise_path = tmp_path / "initial_latents.pt"
spec = dict(model=str(model.resolve()), reference=str(reference.resolve()), prompt=prompt, refs=refs,
resolution=resolution, steps=int(os.environ.get("QWEN_IMAGE21_PARITY_STEPS", "4")),
negative_prompt="" if case == "true_cfg" else None,
true_cfg_scale=2 if case == "true_cfg" else 1, cache=case != "uncached", latents=str(noise_path),
seed=21, reference_resolution=1024, height=resolution, width=resolution)
spec.update(_case_overrides(case))
height, width = spec["height"], spec["width"]
assert min(height, width, spec["reference_resolution"]) > 0
assert all(value % 32 == 0 for value in (height, width, spec["reference_resolution"]))
assert spec["steps"] > 0
spec.setdefault("sigmas", np.linspace(1.0, 1 / spec["steps"], spec["steps"]).tolist())
torch.save(torch.randn(1, (height // 16) * (width // 16), 64,
generator=torch.Generator(device="cpu").manual_seed(spec["seed"]),
dtype=torch.bfloat16), noise_path)
spec_path = tmp_path / "spec.json"
spec_path.write_text(json.dumps(spec))
outputs = []
for kind in ("native", "reference"):
output = tmp_path / f"{kind}.pt"
subprocess.run([sys.executable, str(Path(__file__).resolve()), "--worker", kind, str(spec_path), str(output)],
check=True, cwd=ROOT)
outputs.append(torch.load(output, map_location="cpu", weights_only=True))
actual, expected = outputs[0]["pixels"], outputs[1]["pixels"]
assert actual.shape == expected.shape == (1, 4, 1, height, width)
assert torch.isfinite(actual).all() and bool(((actual >= 0) & (actual <= 1)).all())
assert torch.isfinite(expected).all() and bool(((expected >= 0) & (expected <= 1)).all())
metrics = {"case": case, "steps": spec["steps"], "cache": spec["cache"], "seed": spec["seed"],
"rgb": _difference_metrics(actual[:, :3], expected[:, :3]),
"all_channels": _difference_metrics(actual, expected)}
(tmp_path / "parity_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
print(json.dumps(metrics), flush=True)
# Provisional tolerances: calibrate from measured BF16 CUDA evidence before release.
torch.testing.assert_close(actual, expected, atol=0.04, rtol=0.02)
native_png = Image.open(outputs[0]["png_path"])
assert native_png.mode == "RGBA"
expected_u8 = (actual[0, :, 0].permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8).numpy()
np.testing.assert_array_equal(np.array(native_png), expected_u8)
if outputs[0]["peak_memory_mb"] is not None:
assert outputs[0]["peak_memory_mb"] <= 24 * 1024
@pytest.mark.parametrize("case", CASES)
def test_qwen_image21_modes_1024_pipeline_parity(case, tmp_path):
_run_parity(case, 1024, tmp_path)
def test_qwen_image21_t2i_2k_pipeline_parity(tmp_path):
if os.environ.get("QWEN_IMAGE21_RUN_2K_PARITY") != "1":
pytest.skip("Set QWEN_IMAGE21_RUN_2K_PARITY=1 for the separate 2K gate")
if os.environ.get("QWEN_IMAGE21_RUN_PIPELINE_PARITY") != "1":
pytest.fail("The 2K gate also requires QWEN_IMAGE21_RUN_PIPELINE_PARITY=1", pytrace=False)
_run_parity("t2i", 2048, tmp_path)
if __name__ == "__main__":
if len(sys.argv) != 5 or sys.argv[1] != "--worker":
raise SystemExit("This helper is launched by pytest")
_worker(sys.argv[2], sys.argv[3], sys.argv[4])
tests/local_tests/transformers/test_qwen_image21_transformer_parity.py
+268
View File
@@ -0,0 +1,268 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 transformer reference and production-loader parity.
Set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers revision below for random-weight
CPU parity. Real-weight BF16 parity additionally needs an allocated CUDA GPU,
QWEN_IMAGE21_MODEL_ROOT, and QWEN_IMAGE21_RUN_DIT_PARITY=1. Explicitly requested
real-weight tests fail rather than skip when their prerequisites are missing.
Set QWEN_IMAGE21_DIT_PARITY_REPORT to save block and prediction diagnostics as JSON.
"""
from __future__ import annotations
import gc
import importlib
import json
import math
import os
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
from torch.testing import assert_close
ROOT = Path(__file__).resolve().parents[3]
REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
REFERENCE_ROOT = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", ROOT / "official_reference/diffusers"))
MODEL_ROOT = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", ROOT / "official_weights/Qwen-Image-2.1"))
RUN_ENV = "QWEN_IMAGE21_RUN_DIT_PARITY"
PARITY_SCOPE = "both"
def _numerical_metrics(actual, expected):
actual, expected = actual.detach().float().cpu(), expected.detach().float().cpu()
difference = actual - expected
reference_rms = expected.square().mean().sqrt()
reference_abs_mean = expected.abs().mean()
rmse = difference.square().mean().sqrt()
mean_abs = difference.abs().mean()
floor = torch.finfo(torch.float32).tiny
return {
"shape": list(actual.shape),
"finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()),
"max_abs": difference.abs().max().item(),
"mean_abs": mean_abs.item(),
"rmse": rmse.item(),
"reference_rms": reference_rms.item(),
"reference_abs_mean": reference_abs_mean.item(),
"relative_l2": (rmse / reference_rms.clamp_min(floor)).item(),
"relative_abs_mean": (mean_abs / reference_abs_mean.clamp_min(floor)).item(),
"total_elements": actual.numel(),
}
def _modality_summaries(actual, expected, masks):
actual, expected = actual.detach().cpu(), expected.detach().cpu()
summaries = {"global": _numerical_metrics(actual, expected)}
for name, mask in masks.items():
assert mask.numel() == actual.shape[1], f"{name} token mask does not match the joint sequence"
if bool(mask.any()):
summaries[name] = _numerical_metrics(actual[:, mask], expected[:, mask])
return summaries
def _reference_module(required=False):
if not (REFERENCE_ROOT / "src").is_dir():
message = f"Set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers checkout at {REFERENCE_REVISION}"
if required:
pytest.fail(message, pytrace=False)
pytest.skip(message)
revision = subprocess.check_output(["git", "-C", str(REFERENCE_ROOT), "rev-parse", "HEAD"], text=True).strip()
assert revision == REFERENCE_REVISION, f"Reference revision mismatch: {revision}"
sys.path.insert(0, str(REFERENCE_ROOT / "src"))
try:
module = importlib.import_module("diffusers.models.transformers.transformer_qwenimage21")
except ImportError as exc:
if required:
pytest.fail(f"Qwen-Image-2.1 reference dependencies unavailable: {exc}", pytrace=False)
pytest.skip(f"Qwen-Image-2.1 reference dependencies unavailable: {exc}")
expected = REFERENCE_ROOT / "src/diffusers/models/transformers"
assert Path(module.__file__).resolve().parent == expected.resolve()
return module
def _inputs(channels=4, text_dim=12, device="cpu", dtype=torch.float32, batch=2):
generator = torch.Generator().manual_seed(21)
image = torch.randn(batch, 12, channels, generator=generator).to(device=device, dtype=dtype)
text = torch.randn(batch, 5, text_dim, generator=generator).to(device=device, dtype=dtype)
return dict(hidden_states=image,
encoder_hidden_states=text,
timestep=torch.full((batch,), 0.65, device=device, dtype=dtype),
img_shapes=[[(1, 2, 2), (1, 2, 2), (1, 2, 2)]] * batch,
img_mask=torch.tensor([[False, True, True, False, False, True]] * batch, device=device),
encoder_hidden_states_mask=torch.tensor([[True, True, True, True, False]] * batch, device=device))
@pytest.fixture
def tiny_pair():
reference = _reference_module()
from fastvideo.configs.models.dits.qwen_image21 import QwenImage21ArchConfig, QwenImage21Config
from fastvideo.models.dits.qwen_image21 import QwenImage21Transformer2DModel
from fastvideo.platforms import AttentionBackendEnum
kwargs = dict(in_channels=4, out_channels=4, num_layers=2, attention_head_dim=8,
num_attention_heads=2, context_in_dim=12, mlp_ratio=3, axes_dims_rope=(2, 2, 4))
torch.manual_seed(20)
official = reference.QwenImage21Transformer2DModel(**kwargs).eval()
config = QwenImage21Config(arch_config=QwenImage21ArchConfig(**kwargs))
config._resolved_attention_backend = AttentionBackendEnum.TORCH_SDPA
native = QwenImage21Transformer2DModel(config, {}).eval()
official_state = official.state_dict()
assert {key: value.shape for key, value in native.state_dict().items()} == {
key: value.shape for key, value in official_state.items()
}
native.load_state_dict(official_state, strict=True)
return reference, official, native
def test_qwen_image21_dit_random_reference_parity(tiny_pair):
_, official, native = tiny_pair
inputs = _inputs()
with torch.inference_mode():
expected = official(**inputs, return_dict=False)[0]
actual = native(**inputs)
assert_close(actual, expected, atol=2e-6, rtol=2e-5)
def test_qwen_image21_qk_norm_dtype_reference_parity():
reference = _reference_module()
from fastvideo.models.dits.qwen_image21 import QwenImage21RMSNorm
for weight_dtype in (torch.float32, torch.bfloat16, torch.float16):
for input_dtype in (torch.float32, torch.bfloat16, torch.float16):
official = reference.RMSNorm(8, eps=1e-6).to(weight_dtype)
native = QwenImage21RMSNorm(8, eps=1e-6).to(weight_dtype)
weight = torch.randn(8, generator=torch.Generator().manual_seed(23)).to(weight_dtype)
official.weight.data.copy_(weight)
native.weight.data.copy_(weight)
hidden = torch.randn(2, 7, 3, 8, generator=torch.Generator().manual_seed(24)).to(input_dtype)
assert_close(native(hidden), official(hidden), atol=0, rtol=0)
@pytest.mark.parametrize("storage_device", [None, "cpu"])
def test_qwen_image21_dit_reference_cache_parity(tiny_pair, storage_device):
reference, official, native = tiny_pair
from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache
inputs = _inputs()
official_cache = reference.QwenImage21KVCache(2)
native_cache = QwenImage21KVCache(2, storage_device)
with torch.inference_mode():
expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="extract", return_dict=False)[0]
actual = native(**inputs, kv_cache=native_cache, kv_cache_mode="extract")
assert_close(actual, expected, atol=2e-6, rtol=2e-5)
for layer_idx in range(2):
expected_k, expected_v = official_cache.get_layer(layer_idx).get()
actual_k, actual_v = native_cache.get_layer(layer_idx).get()
assert_close(actual_k, expected_k, atol=2e-6, rtol=2e-5)
assert_close(actual_v, expected_v, atol=2e-6, rtol=2e-5)
inputs["timestep"].fill_(0.3)
inputs["hidden_states"][:, -4:] *= 0.8
expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="cached", return_dict=False)[0]
actual = native(**inputs, kv_cache=native_cache, kv_cache_mode="cached")
assert_close(actual, expected, atol=2e-6, rtol=2e-5)
def test_qwen_image21_dit_production_loader_parity():
if os.environ.get(RUN_ENV) != "1":
pytest.skip(f"Set {RUN_ENV}=1 on an allocated CUDA node")
if not torch.cuda.is_available():
pytest.fail("Qwen-Image-2.1 DiT production-loader parity requires CUDA", pytrace=False)
reference = _reference_module(required=True)
component = MODEL_ROOT / "transformer"
if not (component / "config.json").is_file() or not list(component.glob("*.safetensors")):
pytest.fail(f"Local transformer checkpoint missing at {component}", pytrace=False)
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.models.dits.qwen_image21 import QwenImage21KVCache
from fastvideo.models.loader.component_loader import PipelineComponentLoader
args = SimpleNamespace(pipeline_config=QwenImage21PipelineConfig(), model_paths={},
override_transformer_cls_name=None, hsdp_replicate_dim=1, hsdp_shard_dim=1,
dit_cpu_offload=False, pin_cpu_memory=False, use_fsdp_inference=False,
training_mode=False, enable_torch_compile=False, torch_compile_kwargs={},
inference_torch_compile=False, VSA_tile_size=None, inference_mode=True,
dit_layerwise_offload=False, attention_backend="TORCH_SDPA")
native = PipelineComponentLoader.load_module("transformer", str(component), "diffusers", args)
inputs = _inputs(channels=native.config.in_channels, text_dim=native.config.context_in_dim,
device="cuda:0", dtype=torch.bfloat16, batch=1)
cached_inputs = dict(inputs)
cached_inputs["timestep"] = inputs["timestep"] * 0.5
actual_layers = []
handles = [block.register_forward_hook(lambda module, args, output: actual_layers.append(output.float().cpu()))
for block in native.transformer_blocks]
cache = QwenImage21KVCache(len(native.transformer_blocks), "cpu")
with torch.inference_mode():
actual = native(**inputs, kv_cache=cache, kv_cache_mode="extract").float().cpu()
for handle in handles:
handle.remove()
actual_cached = native(**cached_inputs, kv_cache=cache, kv_cache_mode="cached").float().cpu()
del native, cache
gc.collect()
torch.cuda.empty_cache()
official, info = reference.QwenImage21Transformer2DModel.from_pretrained(component, local_files_only=True,
torch_dtype=torch.bfloat16,
output_loading_info=True)
assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
if info.get(key)}
official = official.to("cuda:0").eval()
expected_layers = []
handles = [block.register_forward_hook(lambda module, args, output: expected_layers.append(output.float().cpu()))
for block in official.transformer_blocks]
official_cache = reference.QwenImage21KVCache(len(official.transformer_blocks))
with torch.inference_mode():
expected = official(**inputs, kv_cache=official_cache, kv_cache_mode="extract", return_dict=False)[0]
for handle in handles:
handle.remove()
expected_cached = official(**cached_inputs, kv_cache=official_cache, kv_cache_mode="cached",
return_dict=False)[0]
# Image slots expand into four unpatched latent tokens; the input masks
# independently identify text, reference images, target image and padding.
slots = inputs["img_mask"][0].bool().cpu()
repeats = torch.where(slots, 4, 1)
image_tokens = slots.repeat_interleave(repeats)
target_tokens = math.prod(inputs["img_shapes"][0][-1])
prefix = torch.arange(image_tokens.numel()) < image_tokens.numel() - target_tokens
prompt_valid = inputs["encoder_hidden_states_mask"][0].bool().cpu()
slot_valid = torch.cat((prompt_valid, torch.ones(slots.numel() - prompt_valid.numel(), dtype=torch.bool)))
text_valid = slot_valid.repeat_interleave(repeats)
masks = {
"text_prefix": ~image_tokens & text_valid & prefix,
"image_tokens": image_tokens,
"reference_images": image_tokens & prefix,
"target_image": image_tokens & ~prefix,
"padding": ~image_tokens & ~text_valid & prefix,
}
report = {
"component": "qwen_image21_transformer",
"checkpoint_revision": MODEL_ROOT.name,
"reference_revision": REFERENCE_REVISION,
"torch_version": torch.__version__,
"atol": 3e-3,
"rtol": 1e-2,
"blocks": [],
}
for index, (actual_layer, expected_layer) in enumerate(zip(actual_layers, expected_layers, strict=True)):
summary = {"block": index, **_modality_summaries(actual_layer, expected_layer, masks)}
report["blocks"].append(summary)
print(f"QWEN_IMAGE21_DIT_PARITY block_{index} {json.dumps(summary, sort_keys=True)}", flush=True)
report["final_prediction"] = _modality_summaries(actual, expected, masks)
report["cached_prediction"] = _modality_summaries(
actual_cached, expected_cached, {"target_image": torch.ones(target_tokens, dtype=torch.bool)})
for name in ("final_prediction", "cached_prediction"):
print(f"QWEN_IMAGE21_DIT_PARITY {name} {json.dumps(report[name], sort_keys=True)}", flush=True)
report_path = os.environ.get("QWEN_IMAGE21_DIT_PARITY_REPORT")
if report_path:
destination = Path(report_path)
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
for index, (actual_layer, expected_layer) in enumerate(zip(actual_layers, expected_layers, strict=True)):
assert_close(actual_layer, expected_layer, atol=3e-3, rtol=1e-2, msg=f"Transformer block {index}")
assert_close(actual, expected.float().cpu(), atol=3e-3, rtol=1e-2)
assert_close(actual_cached, expected_cached.float().cpu(), atol=3e-3, rtol=1e-2)
tests/local_tests/vaes/test_qwen_image21_vae_parity.py
+246
View File
@@ -0,0 +1,246 @@
# SPDX-License-Identifier: Apache-2.0
"""Qwen-Image-2.1 VAE parity against a pinned Diffusers checkout.
Random-weight CPU tests cover the image-specific residual paths and tiling.
The production-loader gate requires one CUDA GPU and user-provided weights.
"""
from __future__ import annotations
import gc
import importlib
import os
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
from PIL import Image
from torch.testing import assert_close
REPO_ROOT = Path(__file__).resolve().parents[3]
REFERENCE_REVISION = "578c9b2c6636ab2424a0e56186268b83623656b2"
REFERENCE_ROOT = Path(os.environ.get("QWEN_IMAGE21_DIFFUSERS_DIR", REPO_ROOT / "official_reference" / "diffusers"))
MODEL_ROOT = Path(os.environ.get("QWEN_IMAGE21_MODEL_ROOT", REPO_ROOT / "official_weights" / "Qwen-Image-2.1"))
RUN_ENV = "QWEN_IMAGE21_RUN_VAE_PARITY"
PARITY_SCOPE = "both"
def _reference_class(required: bool = False):
if not (REFERENCE_ROOT / "src").is_dir():
message = f"set QWEN_IMAGE21_DIFFUSERS_DIR to Diffusers checkout at {REFERENCE_REVISION}"
if required:
pytest.fail(message, pytrace=False)
pytest.skip(message)
revision = subprocess.check_output(["git", "-C", str(REFERENCE_ROOT), "rev-parse", "HEAD"], text=True).strip()
assert revision == REFERENCE_REVISION, f"reference revision mismatch: {revision}"
sys.path.insert(0, str(REFERENCE_ROOT / "src"))
try:
module = importlib.import_module("diffusers.models.autoencoders.autoencoder_kl_qwenimage21")
except ImportError as exc:
if required:
pytest.fail(f"cannot import Qwen-Image-2.1 reference: {exc}", pytrace=False)
pytest.skip(f"reference dependencies unavailable: {exc}")
expected_source = REFERENCE_ROOT / "src" / "diffusers" / "models" / "autoencoders"
assert Path(module.__file__).resolve().parent == expected_source.resolve()
return module.AutoencoderKLQwenImage21
def _tiny_config():
from fastvideo.configs.models.vaes.qwen_image21 import QwenImage21VAEArchConfig, QwenImage21VAEConfig
return QwenImage21VAEConfig(
arch_config=QwenImage21VAEArchConfig(
base_dim=4,
decoder_base_dim=4,
z_dim=4,
dim_mult=(1, 2, 4, 8, 8),
num_res_blocks=1,
latents_mean=(0.0,) * 4,
latents_std=(1.0,) * 4,
),
use_tiling=False,
)
@pytest.fixture
def _tiny_pair():
official_cls = _reference_class()
from fastvideo.models.vaes.qwen_image21 import AutoencoderKLQwenImage21
config = _tiny_config()
arch = config.arch_config
official = official_cls(
base_dim=arch.base_dim,
decoder_base_dim=arch.decoder_base_dim,
z_dim=arch.z_dim,
dim_mult=list(arch.dim_mult),
num_res_blocks=arch.num_res_blocks,
temperal_downsample=list(arch.temperal_downsample),
latents_mean=list(arch.latents_mean),
latents_std=list(arch.latents_std),
).eval()
native = AutoencoderKLQwenImage21(config).eval()
official_state = official.state_dict()
assert {key: value.shape for key, value in native.state_dict().items()} == {
key: value.shape for key, value in official_state.items()
}
native.load_state_dict(official_state, strict=True)
return official, native
@pytest.mark.parametrize("tiled", [False, True])
def test_qwen_image21_vae_random_cpu_parity(_tiny_pair, tiled):
official, native = _tiny_pair
if tiled:
for model in (official, native):
model.enable_tiling(
tile_sample_min_height=32,
tile_sample_min_width=32,
tile_sample_stride_height=16,
tile_sample_stride_width=16,
)
generator = torch.Generator().manual_seed(21)
pixels = torch.rand((2, 4, 1, 48, 64), generator=generator).mul_(2).sub_(1)
with torch.inference_mode():
expected = official.encode(pixels).latent_dist
actual = native.encode(pixels).latent_dist
assert_close(actual.mean, expected.mean, atol=2e-6, rtol=2e-5)
assert_close(actual.logvar, expected.logvar, atol=2e-6, rtol=2e-5)
expected_pixels = official.decode(expected.mode()).sample
actual_pixels = native.decode(expected.mode()).sample
assert_close(actual_pixels, expected_pixels, atol=2e-6, rtol=2e-5)
assert actual_pixels.shape == pixels.shape
def test_qwen_image21_vae_slicing_and_output_contract(_tiny_pair):
official, native = _tiny_pair
for model in (official, native):
model.enable_slicing()
pixels = torch.rand((2, 4, 1, 32, 48), generator=torch.Generator().manual_seed(22)).mul_(2).sub_(1)
with torch.inference_mode():
expected = official.encode(pixels, return_dict=False)[0].mode()
actual = native.encode(pixels, return_dict=False)[0].mode()
assert_close(actual, expected, atol=2e-6, rtol=2e-5)
assert_close(
native.decode(actual, return_dict=False)[0],
official.decode(expected, return_dict=False)[0],
atol=2e-6,
rtol=2e-5,
)
def test_qwen_image21_vae_production_loader_parity():
if os.environ.get(RUN_ENV) != "1":
pytest.skip(f"set {RUN_ENV}=1 on an allocated CUDA node")
if not torch.cuda.is_available():
pytest.fail("Qwen-Image-2.1 VAE production parity requires CUDA", pytrace=False)
official_cls = _reference_class(required=True)
component_path = MODEL_ROOT / "vae"
if not (component_path / "config.json").is_file() or not list(component_path.glob("*.safetensors")):
pytest.fail(f"local VAE checkpoint missing at {component_path}", pytrace=False)
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.models.loader.component_loader import VAELoader
device = torch.device("cuda:0")
args = SimpleNamespace(
pipeline_config=QwenImage21PipelineConfig(),
model_paths={},
vae_cpu_offload=False,
)
args.pipeline_config.vae_precision = "fp32"
args.pipeline_config.vae_config.use_tiling = False
native = VAELoader().load(str(component_path), args)
assert all(parameter.dtype == torch.float32 for parameter in native.parameters())
generator = torch.Generator().manual_seed(23)
pixels = torch.rand((1, 4, 1, 64, 80), generator=generator).mul_(2).sub_(1).to(device)
with torch.inference_mode():
posterior = native.encode(pixels).latent_dist
actual_mean = posterior.mean.cpu()
actual_logvar = posterior.logvar.cpu()
actual_pixels = native.decode(posterior.mode()).sample.cpu()
del native, posterior
gc.collect()
torch.cuda.empty_cache()
official, info = official_cls.from_pretrained(
component_path,
local_files_only=True,
torch_dtype=torch.float32,
output_loading_info=True,
)
assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
if info.get(key)}
official = official.to(device).eval()
with torch.inference_mode():
expected = official.encode(pixels).latent_dist
assert_close(actual_mean, expected.mean.cpu(), atol=1e-5, rtol=1e-4)
assert_close(actual_logvar, expected.logvar.cpu(), atol=1e-5, rtol=1e-4)
assert_close(actual_pixels, official.decode(expected.mode()).sample.cpu(), atol=1e-5, rtol=1e-4)
@pytest.mark.parametrize("height,width,tiled", [(64, 80, False), (288, 320, True)])
def test_qwen_image21_vae_bf16_reference_layout_parity(height, width, tiled):
"""Independent preprocessing must preserve singleton strides for BF16 kernels."""
if os.environ.get(RUN_ENV) != "1":
pytest.skip(f"set {RUN_ENV}=1 on an allocated CUDA node")
if not torch.cuda.is_available():
pytest.fail("Qwen-Image-2.1 BF16 VAE production parity requires CUDA", pytrace=False)
official_cls = _reference_class(required=True)
component_path = MODEL_ROOT / "vae"
if not (component_path / "config.json").is_file() or not list(component_path.glob("*.safetensors")):
pytest.fail(f"local VAE checkpoint missing at {component_path}", pytrace=False)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import QwenImage21Pipeline
from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig
from fastvideo.models.loader.component_loader import VAELoader
from fastvideo.pipelines.basic.qwen_image21.stages import normalize_latents, pack_latents, reference_pixels
rgb = torch.randint(0, 256, (height, width, 3), generator=torch.Generator().manual_seed(23), dtype=torch.uint8)
image = Image.fromarray(rgb.numpy()).convert("RGBA")
processor = VaeImageProcessor(vae_scale_factor=16, do_convert_rgb=False)
native_pixels = reference_pixels(image)
official_pixels = processor.preprocess(image, height=height, width=width).unsqueeze(2)
assert_close(native_pixels, official_pixels, atol=0, rtol=0)
# Size-one strides can select a different BF16 convolution kernel even
# though both tensors contain identical values and are channels_last_3d.
assert native_pixels.stride() == official_pixels.stride()
device = torch.device("cuda:0")
native_pixels = native_pixels.to(device=device, dtype=torch.bfloat16)
official_pixels = official_pixels.to(device=device, dtype=torch.bfloat16)
assert native_pixels.stride() == official_pixels.stride()
args = SimpleNamespace(pipeline_config=QwenImage21PipelineConfig(), model_paths={}, vae_cpu_offload=False)
args.pipeline_config.vae_config.use_tiling = tiled
native = VAELoader().load(str(component_path), args)
assert all(parameter.dtype == torch.bfloat16 for parameter in native.parameters())
with torch.no_grad():
posterior = native.encode(native_pixels).latent_dist
actual_mean, actual_logvar = posterior.mean.cpu(), posterior.logvar.cpu()
normalized = normalize_latents(posterior.mode(), native.config.latents_mean, native.config.latents_std)
actual_packed = pack_latents(normalized).cpu()
del native, posterior, normalized
gc.collect()
torch.cuda.empty_cache()
official, info = official_cls.from_pretrained(
component_path, local_files_only=True, torch_dtype=torch.bfloat16, output_loading_info=True)
assert not {key: info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
if info.get(key)}
official = official.to(device).eval()
if tiled:
official.enable_tiling()
with torch.no_grad():
expected = official.encode(official_pixels).latent_dist
assert_close(actual_mean, expected.mean.cpu(), atol=1e-5, rtol=1e-4)
assert_close(actual_logvar, expected.logvar.cpu(), atol=1e-5, rtol=1e-4)
mean = torch.tensor(official.config.latents_mean).view(1, official.config.z_dim, 1, 1, 1)
std = torch.tensor(official.config.latents_std).view(1, official.config.z_dim, 1, 1, 1)
normalized = (expected.mode() - mean.to(device, torch.bfloat16)) / std.to(device, torch.bfloat16)
expected_packed = QwenImage21Pipeline._pack_latents(
normalized, 1, official.config.z_dim, height // 16, width // 16).cpu()
assert_close(actual_packed, expected_packed, atol=1e-5, rtol=1e-4)