v6e-16, 720p

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