[tests] Improve ideogram4 tests (#13862)

* improve ideogram4 tests

* fix

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Author: sayakpaul

src/diffusers/models/transformers/transformer_ideogram4.py

@@ -136,6 +136,7 @@ def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-5) -> None:
         self.hidden_size = hidden_size
         self.num_heads = num_heads
         self.head_dim = hidden_size // num_heads
+        self.use_bias = False
 
         self.to_q = nn.Linear(hidden_size, hidden_size, bias=False)
         self.to_k = nn.Linear(hidden_size, hidden_size, bias=False)

tests/models/transformers/test_models_transformer_ideogram4.py

@@ -13,8 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
+import pytest
 import torch
 
 from diffusers import Ideogram4Transformer2DModel
@@ -23,19 +22,22 @@
     LLM_TOKEN_INDICATOR,
     OUTPUT_IMAGE_INDICATOR,
 )
+from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..test_modeling_common import ModelTesterMixin
+from ..testing_utils import (
+    AttentionTesterMixin,
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TrainingTesterMixin,
+)
 
 
 enable_full_determinism()
 
 
-class Ideogram4TransformerTests(ModelTesterMixin, unittest.TestCase):
-    model_class = Ideogram4Transformer2DModel
-    main_input_name = "hidden_states"
-    model_split_percents = [0.9, 0.9, 0.9]
-
+class Ideogram4TransformerTesterConfig(BaseModelTesterConfig):
     _hidden_size = 32
     _num_heads = 4
     _head_dim = _hidden_size // _num_heads  # 8
@@ -44,61 +46,32 @@ class Ideogram4TransformerTests(ModelTesterMixin, unittest.TestCase):
     _max_text_tokens = 4
     _num_image_tokens = 4
 
-    def prepare_dummy_input(self, height: int = 0, width: int = 0):
-        del height, width
-        batch_size = 1
-        max_text_tokens = self._max_text_tokens
-        num_image_tokens = self._num_image_tokens
-        seq_len = max_text_tokens + num_image_tokens
-
-        hidden_states = torch.zeros(batch_size, seq_len, self._in_channels)
-        hidden_states[:, max_text_tokens:] = torch.randn(batch_size, num_image_tokens, self._in_channels)
-
-        encoder_hidden_states = torch.zeros(batch_size, seq_len, self._llm_features_dim)
-        encoder_hidden_states[:, :max_text_tokens] = torch.randn(batch_size, max_text_tokens, self._llm_features_dim)
+    @property
+    def model_class(self):
+        return Ideogram4Transformer2DModel
 
-        position_ids = torch.zeros(batch_size, seq_len, 3, dtype=torch.long)
-        text_pos = torch.arange(max_text_tokens)
-        position_ids[:, :max_text_tokens, 0] = text_pos
-        position_ids[:, :max_text_tokens, 1] = text_pos
-        position_ids[:, :max_text_tokens, 2] = text_pos
-        # Image tokens get a 2x2 grid with the IMAGE_POSITION_OFFSET applied.
-        image_h = torch.tensor([0, 0, 1, 1])
-        image_w = torch.tensor([0, 1, 0, 1])
-        position_ids[:, max_text_tokens:, 0] = IMAGE_POSITION_OFFSET
-        position_ids[:, max_text_tokens:, 1] = image_h + IMAGE_POSITION_OFFSET
-        position_ids[:, max_text_tokens:, 2] = image_w + IMAGE_POSITION_OFFSET
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
 
-        segment_ids = torch.ones(batch_size, seq_len, dtype=torch.long)
-        indicator = torch.empty(batch_size, seq_len, dtype=torch.long)
-        indicator[:, :max_text_tokens] = LLM_TOKEN_INDICATOR
-        indicator[:, max_text_tokens:] = OUTPUT_IMAGE_INDICATOR
-        timestep = torch.tensor([0.5])
-
-        inputs = {
-            "hidden_states": hidden_states.to(torch_device),
-            "encoder_hidden_states": encoder_hidden_states.to(torch_device),
-            "timestep": timestep.to(torch_device),
-            "position_ids": position_ids.to(torch_device),
-            "segment_ids": segment_ids.to(torch_device),
-            "indicator": indicator.to(torch_device),
-        }
-        return inputs
+    @property
+    def output_shape(self) -> tuple[int, ...]:
+        return (1, self._max_text_tokens + self._num_image_tokens, self._in_channels)
 
     @property
-    def dummy_input(self):
-        return self.prepare_dummy_input()
+    def input_shape(self) -> tuple[int, ...]:
+        return (1, self._max_text_tokens + self._num_image_tokens, self._in_channels)
 
     @property
-    def input_shape(self):
-        return (self._max_text_tokens + self._num_image_tokens, self._in_channels)
+    def model_split_percents(self) -> list:
+        return [0.9, 0.9, 0.9]
 
     @property
-    def output_shape(self):
-        return (self._max_text_tokens + self._num_image_tokens, self._in_channels)
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
 
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = {
+    def get_init_dict(self) -> dict:
+        return {
             "in_channels": self._in_channels,
             "num_layers": 2,
             "attention_head_dim": self._head_dim,
@@ -110,24 +83,84 @@ def prepare_init_args_and_inputs_for_common(self):
             "mrope_section": (2, 1, 1),
             "norm_eps": 1e-5,
         }
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
+
+    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
+        batch_size = 1
+        max_text_tokens = self._max_text_tokens
+        num_image_tokens = self._num_image_tokens
+        seq_len = max_text_tokens + num_image_tokens
+
+        hidden_states = torch.zeros(
+            batch_size, seq_len, self._in_channels, device=torch_device, dtype=self.torch_dtype
+        )
+        hidden_states[:, max_text_tokens:] = randn_tensor(
+            (batch_size, num_image_tokens, self._in_channels),
+            generator=self.generator,
+            device=torch_device,
+            dtype=self.torch_dtype,
+        )
+
+        encoder_hidden_states = torch.zeros(
+            batch_size, seq_len, self._llm_features_dim, device=torch_device, dtype=self.torch_dtype
+        )
+        encoder_hidden_states[:, :max_text_tokens] = randn_tensor(
+            (batch_size, max_text_tokens, self._llm_features_dim),
+            generator=self.generator,
+            device=torch_device,
+            dtype=self.torch_dtype,
+        )
+
+        position_ids = torch.zeros(batch_size, seq_len, 3, dtype=torch.long, device=torch_device)
+        text_pos = torch.arange(max_text_tokens, device=torch_device)
+        position_ids[:, :max_text_tokens, 0] = text_pos
+        position_ids[:, :max_text_tokens, 1] = text_pos
+        position_ids[:, :max_text_tokens, 2] = text_pos
+        # Image tokens get a 2x2 grid with the IMAGE_POSITION_OFFSET applied.
+        image_h = torch.tensor([0, 0, 1, 1], device=torch_device)
+        image_w = torch.tensor([0, 1, 0, 1], device=torch_device)
+        position_ids[:, max_text_tokens:, 0] = IMAGE_POSITION_OFFSET
+        position_ids[:, max_text_tokens:, 1] = image_h + IMAGE_POSITION_OFFSET
+        position_ids[:, max_text_tokens:, 2] = image_w + IMAGE_POSITION_OFFSET
+
+        segment_ids = torch.ones(batch_size, seq_len, dtype=torch.long, device=torch_device)
+        indicator = torch.empty(batch_size, seq_len, dtype=torch.long, device=torch_device)
+        indicator[:, :max_text_tokens] = LLM_TOKEN_INDICATOR
+        indicator[:, max_text_tokens:] = OUTPUT_IMAGE_INDICATOR
+        timestep = torch.tensor([0.5], device=torch_device, dtype=self.torch_dtype)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "timestep": timestep,
+            "position_ids": position_ids,
+            "segment_ids": segment_ids,
+            "indicator": indicator,
+        }
+
+
+class TestIdeogram4Transformer(Ideogram4TransformerTesterConfig, ModelTesterMixin):
+    """Core model tests for Ideogram 4 Transformer."""
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+    def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
+        # Skip: the non-persistent fp32 RoPE inv_freq buffer is truncated to fp16 by the in-memory
+        # .to(dtype) path but kept fp32 by from_pretrained, so the two outputs diverge well beyond any
+        # meaningful tolerance. Dtype preservation is already covered by test_from_save_pretrained_dtype
+        # and test_keep_in_fp32_modules.
+        pytest.skip("Tolerance requirements too high for meaningful test")
+
+
+class TestIdeogram4TransformerMemory(Ideogram4TransformerTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests for Ideogram 4 Transformer."""
+
+
+class TestIdeogram4TransformerTraining(Ideogram4TransformerTesterConfig, TrainingTesterMixin):
+    """Training tests for Ideogram 4 Transformer."""
 
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"Ideogram4Transformer2DModel"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
-    def test_forward_signature(self):
-        # The model's forward takes packed inputs by position; skip the strict signature check used by the mixin.
-        return
-
-    def test_output(self):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
-        model.to(torch_device)
-        model.eval()
-        with torch.no_grad():
-            output = model(**inputs_dict, return_dict=False)[0]
-        expected = (1, self._max_text_tokens + self._num_image_tokens, self._in_channels)
-        self.assertEqual(tuple(output.shape), expected)
-        self.assertEqual(output.dtype, torch.float32)
+
+class TestIdeogram4TransformerAttention(Ideogram4TransformerTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for Ideogram 4 Transformer."""