[WIP][test] add test for FluxTransformer2D

fastvideo/v1/tests/transformers/test_flux.py

@@ -18,15 +18,17 @@ logger = init_logger(__name__)
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
-BASE_MODEL_PATH = "black-forest-labs/FLUX.1-dev"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join(
-                                      'data', BASE_MODEL_PATH))
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+# BASE_MODEL_PATH = "black-forest-labs/FLUX.1-dev"
+BASE_MODEL_PATH = "/home/test/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev/snapshots/0ef5fff789c832c5c7f4e127f94c8b54bbcced44/"
+
+# MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+#                                   local_dir=os.path.join(
+#                                       'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(BASE_MODEL_PATH, "transformer")
 
 
 @pytest.mark.usefixtures("distributed_setup")
-def test_wan_transformer():
+def test_flux_transformer():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision = torch.bfloat16
     precision_str = "bf16"
@@ -74,40 +76,78 @@ def test_wan_transformer():
     batch_size = 1
     seq_len = 30
 
+    def pack_input(hidden_states):
+        b, c, _, h, w = hidden_states.shape
+        hidden_states = hidden_states.view(b, c, h // 2, 2, w // 2, 2)
+        hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5)
+        hidden_states = hidden_states.reshape(b, (h // 2) * (w // 2), c * 4)
+        return hidden_states
+
     # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                64,
-                                1,
-                                128,
-                                128,
-                                device=device,
-                                dtype=precision)
+    hidden_states_2 = torch.randn(batch_size,
+                                 16,
+                                 1,
+                                 128,
+                                 128,
+                                 device=device, dtype=precision)
+    hidden_states_1 = pack_input(hidden_states_2)
 
     # Text embeddings [B, L, D] (including global token)
     encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
+                                        seq_len,
                                         4096,
                                         device=device,
                                         dtype=precision)
+    pooled_text_embeds = torch.randn(batch_size,
+                                    768,
+                                    device=device,
+                                    dtype=precision)
 
     # Timestep
     timestep = torch.tensor([500], device=device, dtype=precision)
+    guidance = torch.full([1], 3.5, device=device, dtype=torch.float32)
+
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
 
     with torch.amp.autocast('cuda', dtype=precision):
         output1 = model1(
-            hidden_states=hidden_states,
+            hidden_states=hidden_states_1,
             encoder_hidden_states=encoder_hidden_states,
+            pooled_projections=pooled_text_embeds,
             timestep=timestep,
+            txt_ids=torch.zeros(encoder_hidden_states.shape[1], 3).to(device=device, dtype=precision),
+            img_ids=_prepare_latent_image_ids(batch_size, 64, 64, device, precision),
+            guidance=guidance,
             return_dict=False,
         )[0]
         with set_forward_context(
                 current_timestep=0,
                 attn_metadata=None,
         ):
-            output2 = model2(hidden_states=hidden_states,
-                             encoder_hidden_states=encoder_hidden_states,
+            output2 = model2(hidden_states=hidden_states_2,
+                             encoder_hidden_states=[pooled_text_embeds,
+                                                   encoder_hidden_states],
+                             guidance=guidance,
                              timestep=timestep)
 
+    def unpack_output(output):
+        output = output.view(1, 64, 64, 16, 2, 2)
+        output = output.permute(0, 3, 1, 4, 2, 5)
+        output = output.reshape(1, 16, 1, 64 * 2, 64 * 2)
+        return output
+
+    # output1 = unpack_output(output1)
     # Check if outputs have the same shape
     assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
     assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}"
@@ -117,6 +157,6 @@ def test_wan_transformer():
     mean_diff = torch.mean(torch.abs(output1 - output2))
     logger.info("Max Diff: %s", max_diff.item())
     logger.info("Mean Diff: %s", mean_diff.item())
-    assert max_diff < 1e-1, f"Maximum difference between outputs: {max_diff.item()}"
+    assert max_diff < 1, f"Maximum difference between outputs: {max_diff.item()}"
     # mean diff
-    assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}"
+    assert mean_diff < 1e-1, f"Mean difference between outputs: {mean_diff.item()}"