Fix get_dummy_wan_inputs for two-expert Wan2.2 I2V pipeline

get_dummy_wan_inputs (used by WanPipeline.quantize_transformer) assumed a single
pipeline.transformer and the single-transformer prepare_latents() signature.
WanPipelineI2V_2_2 (Wan2.2-I2V-A14B) is two-expert: it has
low_noise_transformer / high_noise_transformer (no .transformer) and a different
prepare_latents() signature, so qwix quantization crashed with
AttributeError: 'WanPipelineI2V_2_2' object has no attribute 'transformer'
and then
TypeError: prepare_latents() got an unexpected keyword argument 'vae_scale_factor_temporal'.

When pipeline.transformer is absent, build dummy latents directly in the
(B, C, F, H, W) layout WanModel.__call__ expects, taking num_channels_latents
from an existing expert. The single-transformer path is unchanged.

Validated on Wan2.2-I2V-A14B (v6e): quantization now proceeds past these errors
into qwix.quantize_model.

Author: ThomasNing

src/maxdiffusion/maxdiffusion_utils.py

@@ -358,15 +358,29 @@ def get_dummy_ltx2_inputs(config, pipeline, batch_size):
 
 
 def get_dummy_wan_inputs(config, pipeline, batch_size):
-  latents = pipeline.prepare_latents(
-      batch_size,
-      vae_scale_factor_temporal=pipeline.vae_scale_factor_temporal,
-      vae_scale_factor_spatial=pipeline.vae_scale_factor_spatial,
-      height=config.height,
-      width=config.width,
-      num_frames=config.num_frames,
-      num_channels_latents=pipeline.transformer.config.in_channels,
-  )
+  if getattr(pipeline, "transformer", None) is not None:
+    latents = pipeline.prepare_latents(
+        batch_size,
+        vae_scale_factor_temporal=pipeline.vae_scale_factor_temporal,
+        vae_scale_factor_spatial=pipeline.vae_scale_factor_spatial,
+        height=config.height,
+        width=config.width,
+        num_frames=config.num_frames,
+        num_channels_latents=pipeline.transformer.config.in_channels,
+    )
+  else:
+    # The two-expert Wan2.2 I2V pipeline (WanPipelineI2V_2_2) has no `.transformer`
+    # and a different `prepare_latents()` signature, so build dummy latents directly
+    # in the (B, C, F, H, W) layout WanModel.__call__ expects.
+    transformer = getattr(pipeline, "low_noise_transformer", None) or pipeline.high_noise_transformer
+    num_channels_latents = transformer.config.in_channels
+    num_latent_frames = (config.num_frames - 1) // pipeline.vae_scale_factor_temporal + 1
+    latent_height = config.height // pipeline.vae_scale_factor_spatial
+    latent_width = config.width // pipeline.vae_scale_factor_spatial
+    latents = jax.random.normal(
+        jax.random.key(config.seed),
+        (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width),
+    )
   bsz = latents.shape[0]
   prompt_embeds = jax.random.normal(jax.random.key(config.seed), (batch_size, 512, 4096))
   timesteps = jnp.array([0] * bsz, dtype=jnp.int32)