From e8a407bf6e06ac487a72ac0eac8da698a38782a8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 15:29:34 +0000 Subject: [PATCH] Load Wan VAE in fp32 via per-component dtype dict; fix A14B 720p preset --- studio/backend/core/inference/video.py | 24 ++++++++++++------- .../backend/core/inference/video_families.py | 8 ++++--- studio/backend/tests/test_video_families.py | 5 +++- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 7d7bdf6be6..7ad96b6201 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -838,6 +838,14 @@ class VideoBackend: # ── build the pipeline. pipeline_cls = getattr(diffusers, fam.pipeline_class) pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + if getattr(fam, "vae_force_fp32", False): + # Wan's VAE must decode in float32, but a scalar torch_dtype casts EVERY component + # (VAE included) to the pipe dtype during load -- AutoencoderKLWan has no + # _keep_in_fp32_modules, so from_pretrained truncates its fp32 weights to bf16 and a + # later .to(float32) only widens the already-lossy values (banding / black frames). + # diffusers >= 0.39 takes a per-component dtype dict, so load the VAE at fp32 directly; + # "default" MUST be set or unlisted components fall back to fp32 (over-widening the DiT). + pipe_kwargs["torch_dtype"] = {"vae": torch.float32, "default": dtype} if hf_token: pipe_kwargs["token"] = hf_token if kind == "pipeline": @@ -878,14 +886,14 @@ class VideoBackend: _base_local_dir or base, transformer = transformer, **pipe_kwargs ) - # Wan's VAE decodes in float32: diffusers loads AutoencoderKLWan at torch.float32 while - # the pipe runs bf16 (WanPipeline docstring). from_pretrained above cast every component - # (VAE included) to the pipe dtype, so pin the VAE back to fp32 or every clip decodes with - # banding / black frames. bf16_components_gb already budgets the VAE at its fp32 size, so - # the memory plan stays consistent. - if getattr(fam, "vae_force_fp32", False) and getattr(pipe, "vae", None) is not None: - import torch - pipe.vae.to(torch.float32) + # The per-component torch_dtype above already loads the Wan VAE at float32 (bf16_components_gb + # budgets it at that fp32 size, so the memory plan stays consistent). Belt-and-suspenders for + # any path that bypassed the dict (e.g. a passed-in vae=): re-pin an fp32-force VAE that came + # back at a lower precision. This is a no-op on the primary path (the load already fp32'd it). + if getattr(fam, "vae_force_fp32", False): + vae = getattr(pipe, "vae", None) + if vae is not None and getattr(vae, "dtype", None) is not torch.float32: + vae.to(torch.float32) if _load_token is not None and _load_token != self._load_token: del pipe diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index 9feccf5100..835b968bc3 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -191,9 +191,11 @@ _FAMILIES: tuple[VideoFamily, ...] = ( default_fps = 16, frame_step = 4, resolution_multiple = 16, - # 480p and 720p presets (landscape, vertical, square), the two resolutions the - # A14B card documents. 832x480 is the native 480p; 1280x704 the 720p target. - resolution_presets = ((1280, 704), (832, 480), (480, 832), (704, 1280)), + # 480p and 720p presets (landscape + vertical), the two resolutions the A14B card + # documents. 832x480 is the native 480p; 1280x720 the native 720p (true 16:9). A14B's + # VAE is 8x so resolution_multiple is 16 and 720 (= 45*16) renders exactly -- the 704 + # value belongs to TI2V-5B, whose 16x VAE floors 720 to 704 (multiple 32). + resolution_presets = ((1280, 720), (832, 480), (480, 832), (720, 1280)), # bf16-RESIDENT sizes. Each expert ships FP32 on disk (safetensors headers are F32; # transformer index = 57.15 GB = 14.3B params x 4), so bf16-resident is ~28.6 each -> # ~57.2 for BOTH experts (the memory headline before offload), NOT the 114.3 fp32 diff --git a/studio/backend/tests/test_video_families.py b/studio/backend/tests/test_video_families.py index b9ee682209..fa1ef242e9 100644 --- a/studio/backend/tests/test_video_families.py +++ b/studio/backend/tests/test_video_families.py @@ -164,7 +164,10 @@ def test_wan_snap_video_size_16(): # Wan patchifies at spatial factor 8 * patch 2 = 16; sizes floor to /16. fam = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers") assert fam.resolution_multiple == 16 - assert snap_video_size(fam, 1280, 704) == (1280, 704) # on-grid preset + # A14B's native 720p is the true 16:9 1280x720 (720 = 45*16 renders exactly on the /16 grid), + # NOT the 1280x704 that TI2V-5B's /32 VAE floors to. The default preset is that native 720p. + assert fam.resolution_presets[0] == (1280, 720) + assert snap_video_size(fam, 1280, 720) == (1280, 720) # native 720p, on-grid assert snap_video_size(fam, 1000, 700) == (992, 688)