testing 480p

Author: prishajain1

src/maxdiffusion/tests/wan_magcache_test.py

@@ -75,8 +75,8 @@ def setUpClass(cls):
             None,
             os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
             "num_inference_steps=50",
-            "height=720",
-            "width=1280",
+            "height=480",
+            "width=832",
             "num_frames=81",
             "seed=11234567893",
             "attention=flash",
@@ -148,8 +148,8 @@ def setUpClass(cls):
             None,
             os.path.join(THIS_DIR, "..", "configs", "base_wan_27b.yml"),
             "num_inference_steps=40",
-            "height=720",
-            "width=1280",
+            "height=480",
+            "width=832",
             "num_frames=81",
             "guidance_scale_low=3.0",
             "guidance_scale_high=4.0",
@@ -161,7 +161,7 @@ def setUpClass(cls):
             "skip_jax_distributed_system=True",
             "weights_dtype=bfloat16",
             "activations_dtype=bfloat16",
-            "per_device_batch_size=0.125",
+            "per_device_batch_size=0.25",
             "ici_data_parallelism=2",
             "ici_fsdp_parallelism=1",
             "ici_context_parallelism=4",
@@ -230,8 +230,9 @@ def setUpClass(cls):
             None,
             os.path.join(THIS_DIR, "..", "configs", "base_wan_i2v_14b.yml"),
             "num_inference_steps=50",
-            "height=720",
-            "width=1280",
+            "height=480",
+            "width=832",
+            "pretrained_model_name_or_path=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
             "num_frames=81",
             "seed=11234567893",
             "attention=flash",
@@ -308,8 +309,8 @@ def setUpClass(cls):
             None,
             os.path.join(THIS_DIR, "..", "configs", "base_wan_i2v_27b.yml"),
             "num_inference_steps=40",
-            "height=720",
-            "width=1280",
+            "height=480",
+            "width=832",
             "num_frames=81",
             "guidance_scale_low=3.0",
             "guidance_scale_high=4.0",
@@ -321,7 +322,7 @@ def setUpClass(cls):
             "skip_jax_distributed_system=True",
             "weights_dtype=bfloat16",
             "activations_dtype=bfloat16",
-            "per_device_batch_size=0.125",
+            "per_device_batch_size=0.25",
             "ici_data_parallelism=2",
             "ici_fsdp_parallelism=1",
             "ici_context_parallelism=4",