Load Wan VAE in fp32 via per-component dtype dict; fix A14B 720p preset

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

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

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)