Enable VAE tiling for Qwen-Image-2.1 (#5)

* init

Signed-off-by: Isotr0py <mozf@inferact.ai>

* fix

Signed-off-by: Isotr0py <mozf@inferact.ai>

* test(diffusion): add Qwen-Image-2.1 unit tests

Signed-off-by: Isotr0py <mozf@inferact.ai>

* fix(diffusion): Qwen-Image-2.1 VAE spatial ratio, tiling e2e verification

- AutoencoderKLQwenImage21: derive spatial_compression_ratio from the
  architecture (2 ** len(temperal_downsample) = 16x) instead of trusting
  the checkpoint's scale_factor_spatial=8, which corrupted tiled
  decode/encode geometry (output cropped to half size) and tiling gating.
- pipeline: _read_vae_scale_factor prefers the architecture-derived value
  for the same reason.
- models/qwen_image_21/__init__: lazy exports to break the circular import
  distributed VAE -> package __init__ -> pipeline -> distributed VAE.
- Add VAE unit tests: ratio derivation, tiled vs untiled decode/encode
  consistency, distributed tile_split/exec/merge == tiled_decode.
- e2e matrix (GB200, seed 42, 20 steps) all green; see e2e_out/qwen21_vae/.
- examples: qwen_image_21_edit.py gains --vae-use-tiling.
- docs: mark VAE patch-parallel verified for Qwen-Image-2.1; recipe notes
  the default-tile banding limitation at 2048px+.

* clean tests

Signed-off-by: Isotr0py <mozf@inferact.ai>

* fix cudnn determination

Signed-off-by: Isotr0py <mozf@inferact.ai>

---------

Signed-off-by: Isotr0py <mozf@inferact.ai>
This commit is contained in:
Isotr0py authored and GitHub committed 2026-09-18 10:02:50 +08:00
1 parent 4fc3cf8c60
commit 3affc7766d
7 files changed
+163 -29

No files matched your search

@@ -144,6 +144,11 @@ def parse_args() -> argparse.Namespace:
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--vae-use-tiling",
action="store_true",
help="Enable VAE tiling for memory optimization (also tiles condition-image encoding).",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
@@ -174,6 +179,7 @@ def main():
"model": args.model,
"ulysses_degree": args.ulysses_degree,
"cfg_parallel_size": args.cfg_parallel_size,
"vae_use_tiling": args.vae_use_tiling,
}
if args.enforce_eager is not None:
omni_kwargs["enforce_eager"] = args.enforce_eager