reinitialize the padding tokens to ones to prevent NaN problems.

Author: sayakpaul

tests/lora/test_lora_layers_z_image.py

@@ -37,10 +37,10 @@
 from .utils import PeftLoraLoaderMixinTests  # noqa: E402
 
 
-@unittest.skip(
-    "ZImage LoRA tests are skipped due to non-deterministic behavior from complex64 RoPE operations "
-    "and torch.empty padding tokens. LoRA functionality works correctly with real models."
-)
+# @unittest.skip(
+#     "ZImage LoRA tests are skipped due to non-deterministic behavior from complex64 RoPE operations "
+#     "and torch.empty padding tokens. LoRA functionality works correctly with real models."
+# )
 @require_peft_backend
 class ZImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
     pipeline_class = ZImagePipeline
@@ -127,6 +127,11 @@ def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=No
         tokenizer = Qwen2Tokenizer.from_pretrained(self.tokenizer_id)
 
         transformer = self.transformer_cls(**self.transformer_kwargs)
+        # `x_pad_token` and `cap_pad_token` are initialized with `torch.empty`.
+        # This can cause NaN data values in our testing environment. Fixating them
+        # helps prevent that issue.
+        transformer.x_pad_token.copy_(torch.ones_like(transformer.x_pad_token.data))
+        transformer.cap_pad_token.copy_(torch.ones_like(transformer.cap_pad_token.data))
         vae = self.vae_cls(**self.vae_kwargs)
 
         if scheduler_cls is None:

tests/pipelines/z_image/test_z_image.py

@@ -101,6 +101,11 @@ def get_dummy_components(self):
             axes_dims=[8, 4, 4],
             axes_lens=[256, 32, 32],
         )
+        # `x_pad_token` and `cap_pad_token` are initialized with `torch.empty`.
+        # This can cause NaN data values in our testing environment. Fixating them
+        # helps prevent that issue.
+        transformer.x_pad_token.copy_(torch.ones_like(transformer.x_pad_token.data))
+        transformer.cap_pad_token.copy_(torch.ones_like(transformer.cap_pad_token.data))
 
         torch.manual_seed(0)
         vae = AutoencoderKL(