Disable CFG cache on wan 2.1 when batch size is indivisible by data shards

Author: Perseus14

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -226,6 +226,18 @@ def run_inference_2_1(
   do_cfg = guidance_scale > 1.0
   bsz = latents.shape[0]
 
+  data_shards = 1
+  try:
+    if hasattr(latents, "sharding") and hasattr(latents.sharding, "mesh"):
+      data_shards = latents.sharding.mesh.shape["data"] * latents.sharding.mesh.shape.get("fsdp", 1)
+  except Exception:
+    pass
+
+  if use_cfg_cache and do_cfg and bsz % data_shards != 0:
+    from maxdiffusion.max_logging import MaxLogging
+    max_logging = MaxLogging()
+    max_logging.log(f"Warning: Disabling CFG cache because batch size {bsz} is not divisible by data shards {data_shards}. This often happens with data_parallelism > 1 and per_device_batch_size = 1.")
+    use_cfg_cache = False
   # Resolution-dependent CFG cache config (FasterCache / MixCache guidance)
   if height >= 720:
     # 720p: conservative — protect last 40%, interval=5