diff --git a/.github/scripts/plan_merge_ci.py b/.github/scripts/plan_merge_ci.py index 10524392b..e9faa4019 100644 --- a/.github/scripts/plan_merge_ci.py +++ b/.github/scripts/plan_merge_ci.py @@ -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", ), diff --git a/docs-site/lib/cookbook-pages.ts b/docs-site/lib/cookbook-pages.ts index b44799a35..3dac83a84 100644 --- a/docs-site/lib/cookbook-pages.ts +++ b/docs-site/lib/cookbook-pages.ts @@ -111,6 +111,7 @@ const familyInfo: Record = [ [/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", }; diff --git a/docs-site/navigation.json b/docs-site/navigation.json index d0c60a558..db8f15868 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -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" diff --git a/docs-site/tests/cookbook.spec.ts b/docs-site/tests/cookbook.spec.ts index 14188f0a3..0fb940950 100644 --- a/docs-site/tests/cookbook.spec.ts +++ b/docs-site/tests/cookbook.spec.ts @@ -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 }) => { diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index ae963efa1..e19feae8a 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -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", diff --git a/docs/assets/logos/SOURCES.md b/docs/assets/logos/SOURCES.md index 17b3907f4..619e476c5 100644 --- a/docs/assets/logos/SOURCES.md +++ b/docs/assets/logos/SOURCES.md @@ -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) | diff --git a/docs/assets/logos/qwen.webp b/docs/assets/logos/qwen.webp new file mode 100644 index 000000000..ff987bfef Binary files /dev/null and b/docs/assets/logos/qwen.webp differ diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md index 46d64141e..10a75432c 100644 --- a/docs/cookbook/index.md +++ b/docs/cookbook/index.md @@ -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) diff --git a/docs/cookbook/qwen-image21.md b/docs/cookbook/qwen-image21.md new file mode 100644 index 000000000..2dd24422e --- /dev/null +++ b/docs/cookbook/qwen-image21.md @@ -0,0 +1,108 @@ +# Qwen-Image-2.1 recipes + + + +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 `. 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. diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 1228b1586..9027bb660 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -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." diff --git a/docs/generate_examples.py b/docs/generate_examples.py index 9dffdf188..b962e8ae8 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -36,6 +36,7 @@ COOKBOOK_FAMILIES = { "kandinsky5", "kandinsky6", "flux", + "qwen_image21", "sd35", "minimax_h3", "stable_audio", diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 80eb390e6..722e5cf28 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -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)
[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)
[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) | diff --git a/examples/inference/basic/basic_qwen_image21.py b/examples/inference/basic/basic_qwen_image21.py new file mode 100644 index 000000000..352ef2968 --- /dev/null +++ b/examples/inference/basic/basic_qwen_image21.py @@ -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() diff --git a/examples/inference/basic/qwen_image21_edit.yaml b/examples/inference/basic/qwen_image21_edit.yaml new file mode 100644 index 000000000..b94ae2137 --- /dev/null +++ b/examples/inference/basic/qwen_image21_edit.yaml @@ -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 diff --git a/examples/inference/basic/qwen_image21_t2i.yaml b/examples/inference/basic/qwen_image21_t2i.yaml new file mode 100644 index 000000000..fc8ed335c --- /dev/null +++ b/examples/inference/basic/qwen_image21_t2i.yaml @@ -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 diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index bba086442..ab4a4e0f7 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -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" ] diff --git a/fastvideo/configs/models/dits/qwen_image21.py b/fastvideo/configs/models/dits/qwen_image21.py new file mode 100644 index 000000000..266cc980e --- /dev/null +++ b/fastvideo/configs/models/dits/qwen_image21.py @@ -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" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 4cea0fb33..d2797f2db 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -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", diff --git a/fastvideo/configs/models/encoders/qwen_image21.py b/fastvideo/configs/models/encoders/qwen_image21.py new file mode 100644 index 000000000..98d21e0c5 --- /dev/null +++ b/fastvideo/configs/models/encoders/qwen_image21.py @@ -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) diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index acc5ca9c0..e86a59af9 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -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", ] diff --git a/fastvideo/configs/models/vaes/qwen_image21.py b/fastvideo/configs/models/vaes/qwen_image21.py new file mode 100644 index 000000000..fba097510 --- /dev/null +++ b/fastvideo/configs/models/vaes/qwen_image21.py @@ -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 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 73636280c..f42e517b8 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -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" ] diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 17c29bae3..7faf283b4 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -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) diff --git a/fastvideo/configs/pipelines/qwen_image21.py b/fastvideo/configs/pipelines/qwen_image21.py new file mode 100644 index 000000000..5fb95c146 --- /dev/null +++ b/fastvideo/configs/pipelines/qwen_image21.py @@ -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" diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 64a40d671..89b9b3953 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -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() diff --git a/fastvideo/models/dits/qwen_image21.py b/fastvideo/models/dits/qwen_image21.py new file mode 100644 index 000000000..9230029b1 --- /dev/null +++ b/fastvideo/models/dits/qwen_image21.py @@ -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 diff --git a/fastvideo/models/encoders/qwen_image21.py b/fastvideo/models/encoders/qwen_image21.py new file mode 100644 index 000000000..62b6926ce --- /dev/null +++ b/fastvideo/models/encoders/qwen_image21.py @@ -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 diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index f5e5d0add..8cea9dde7 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -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: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 0ba35fb97..897e3f276 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -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, diff --git a/fastvideo/models/vaes/qwen_image21.py b/fastvideo/models/vaes/qwen_image21.py new file mode 100644 index 000000000..45ac7a292 --- /dev/null +++ b/fastvideo/models/vaes/qwen_image21.py @@ -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 diff --git a/fastvideo/pipelines/basic/qwen_image21/__init__.py b/fastvideo/pipelines/basic/qwen_image21/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/pipelines/basic/qwen_image21/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/qwen_image21/presets.py b/fastvideo/pipelines/basic/qwen_image21/presets.py new file mode 100644 index 000000000..a1087405d --- /dev/null +++ b/fastvideo/pipelines/basic/qwen_image21/presets.py @@ -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) diff --git a/fastvideo/pipelines/basic/qwen_image21/qwen_image21_pipeline.py b/fastvideo/pipelines/basic/qwen_image21/qwen_image21_pipeline.py new file mode 100644 index 000000000..8d2a3ecc9 --- /dev/null +++ b/fastvideo/pipelines/basic/qwen_image21/qwen_image21_pipeline.py @@ -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 diff --git a/fastvideo/pipelines/basic/qwen_image21/stages.py b/fastvideo/pipelines/basic/qwen_image21/stages.py new file mode 100644 index 000000000..c24cd804f --- /dev/null +++ b/fastvideo/pipelines/basic/qwen_image21/stages.py @@ -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_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 diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 2df3914f0..d5b743ba1 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -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: diff --git a/fastvideo/tests/golden_gate/test_qwen_image21.py b/fastvideo/tests/golden_gate/test_qwen_image21.py new file mode 100644 index 000000000..f32ec4a42 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_qwen_image21.py @@ -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) diff --git a/fastvideo/tests/ssim/test_qwen_image21_similarity.py b/fastvideo/tests/ssim/test_qwen_image21_similarity.py new file mode 100644 index 000000000..ae3b6af65 --- /dev/null +++ b/fastvideo/tests/ssim/test_qwen_image21_similarity.py @@ -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)], + ) diff --git a/fastvideo/tests/stages/test_qwen_image21_stages.py b/fastvideo/tests/stages/test_qwen_image21_stages.py new file mode 100644 index 000000000..d5a9bebb6 --- /dev/null +++ b/fastvideo/tests/stages/test_qwen_image21_stages.py @@ -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 []) diff --git a/fastvideo/tests/transformers/test_qwen_image21_attention.py b/fastvideo/tests/transformers/test_qwen_image21_attention.py new file mode 100644 index 000000000..2fa32f919 --- /dev/null +++ b/fastvideo/tests/transformers/test_qwen_image21_attention.py @@ -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 diff --git a/tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py b/tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py new file mode 100644 index 000000000..a1d7f5f9a --- /dev/null +++ b/tests/local_tests/encoders/test_qwen_image21_conditioner_parity.py @@ -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"<|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) diff --git a/tests/local_tests/qwen_image21/README.md b/tests/local_tests/qwen_image21/README.md new file mode 100644 index 000000000..678b550e7 --- /dev/null +++ b/tests/local_tests/qwen_image21/README.md @@ -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. diff --git a/tests/local_tests/qwen_image21/test_pipeline_parity.py b/tests/local_tests/qwen_image21/test_pipeline_parity.py new file mode 100644 index 000000000..91ea90c7f --- /dev/null +++ b/tests/local_tests/qwen_image21/test_pipeline_parity.py @@ -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]) diff --git a/tests/local_tests/transformers/test_qwen_image21_transformer_parity.py b/tests/local_tests/transformers/test_qwen_image21_transformer_parity.py new file mode 100644 index 000000000..e480c2d37 --- /dev/null +++ b/tests/local_tests/transformers/test_qwen_image21_transformer_parity.py @@ -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) diff --git a/tests/local_tests/vaes/test_qwen_image21_vae_parity.py b/tests/local_tests/vaes/test_qwen_image21_vae_parity.py new file mode 100644 index 000000000..57cebc0f3 --- /dev/null +++ b/tests/local_tests/vaes/test_qwen_image21_vae_parity.py @@ -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)