update

src/diffusers/commands/daggr_app.py

@@ -1,447 +0,0 @@
-# Copyright 2025 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import ast
-import re
-from argparse import ArgumentParser, Namespace
-from collections import OrderedDict
-from dataclasses import dataclass, field
-
-from ..utils import logging
-from . import BaseDiffusersCLICommand
-
-
-logger = logging.get_logger("diffusers-cli/daggr")
-
-INTERNAL_TYPE_NAMES = {
-    "Tensor",
-    "Generator",
-}
-
-INTERNAL_TYPE_FULL_NAMES = {
-    "torch.Tensor",
-    "torch.Generator",
-    "torch.dtype",
-}
-
-SLIDER_PARAMS = {
-    "height": {"minimum": 256, "maximum": 2048, "step": 64},
-    "width": {"minimum": 256, "maximum": 2048, "step": 64},
-    "num_inference_steps": {"minimum": 1, "maximum": 100, "step": 1},
-    "guidance_scale": {"minimum": 0, "maximum": 30, "step": 0.5},
-    "strength": {"minimum": 0, "maximum": 1, "step": 0.05},
-    "control_guidance_start": {"minimum": 0, "maximum": 1, "step": 0.05},
-    "control_guidance_end": {"minimum": 0, "maximum": 1, "step": 0.05},
-    "controlnet_conditioning_scale": {"minimum": 0, "maximum": 2, "step": 0.1},
-}
-
-
-@dataclass
-class BlockInfo:
-    name: str
-    class_name: str
-    description: str
-    inputs: list
-    outputs: list
-    user_inputs: list = field(default_factory=list)
-    port_connections: list = field(default_factory=list)
-    fixed_inputs: list = field(default_factory=list)
-
-
-def daggr_command_factory(args: Namespace):
-    return DaggrCommand(
-        repo_id=args.repo_id,
-        output=args.output or "daggr_app.py",
-        workflow=getattr(args, "workflow", None),
-        trigger_inputs=getattr(args, "trigger_inputs", None),
-    )
-
-
-class DaggrCommand(BaseDiffusersCLICommand):
-    @staticmethod
-    def register_subcommand(parser: ArgumentParser):
-        daggr_parser = parser.add_parser("daggr", help="Generate a daggr app from a modular pipeline repo.")
-        daggr_parser.add_argument(
-            "repo_id",
-            type=str,
-            help="HuggingFace Hub repo ID containing a modular pipeline (with modular_model_index.json).",
-        )
-        daggr_parser.add_argument(
-            "--output",
-            type=str,
-            default="daggr_app.py",
-            help="Output file path for the generated daggr app. Default: daggr_app.py",
-        )
-        daggr_parser.add_argument(
-            "--workflow",
-            type=str,
-            default=None,
-            help="Named workflow to resolve conditional blocks (e.g. 'text2image', 'image2image').",
-        )
-        daggr_parser.add_argument(
-            "--trigger-inputs",
-            nargs="*",
-            default=None,
-            help="Trigger input names for manual conditional resolution.",
-        )
-        daggr_parser.set_defaults(func=daggr_command_factory)
-
-    def __init__(
-        self,
-        repo_id: str,
-        output: str = "daggr_app.py",
-        workflow: str | None = None,
-        trigger_inputs: list | None = None,
-    ):
-        self.repo_id = repo_id
-        self.output = output
-        self.workflow = workflow
-        self.trigger_inputs = trigger_inputs
-
-    def run(self):
-        from ..modular_pipelines.modular_pipeline import ModularPipelineBlocks
-
-        logger.info(f"Loading blocks from {self.repo_id}...")
-        blocks = ModularPipelineBlocks.from_pretrained(self.repo_id, trust_remote_code=True)
-        blocks_class_name = blocks.__class__.__name__
-
-        if self.workflow:
-            logger.info(f"Resolving workflow: {self.workflow}")
-            exec_blocks = blocks.get_workflow(self.workflow)
-        elif self.trigger_inputs:
-            trigger_kwargs = {name: True for name in self.trigger_inputs}
-            logger.info(f"Resolving with trigger inputs: {self.trigger_inputs}")
-            exec_blocks = blocks.get_execution_blocks(**trigger_kwargs)
-        else:
-            logger.info("Resolving default execution blocks...")
-            exec_blocks = blocks.get_execution_blocks()
-
-        block_infos = _analyze_blocks(exec_blocks)
-        _classify_inputs(block_infos)
-
-        workflow_label = self.workflow or "default"
-        workflow_resolve_code = self._get_workflow_resolve_code()
-        code = _generate_code(block_infos, self.repo_id, blocks_class_name, workflow_label, workflow_resolve_code)
-
-        try:
-            ast.parse(code)
-        except SyntaxError as e:
-            logger.warning(f"Generated code has syntax error: {e}")
-
-        with open(self.output, "w") as f:
-            f.write(code)
-
-        logger.info(f"Daggr app written to {self.output}")
-        print(f"Generated daggr app: {self.output}")
-        print(f"  Pipeline: {blocks_class_name}")
-        print(f"  Workflow: {workflow_label}")
-        print(f"  Blocks: {len(block_infos)}")
-        print(f"\nRun with: python {self.output}")
-
-    def _get_workflow_resolve_code(self):
-        if self.workflow:
-            return f"_pipeline._blocks.get_workflow({self.workflow!r})"
-        elif self.trigger_inputs:
-            kwargs_str = ", ".join(f"{name!r}: True" for name in self.trigger_inputs)
-            return f"_pipeline._blocks.get_execution_blocks(**{{{kwargs_str}}})"
-        else:
-            return "_pipeline._blocks.get_execution_blocks()"
-
-
-def _analyze_blocks(exec_blocks):
-    block_infos = []
-    for name, block in exec_blocks.sub_blocks.items():
-        info = BlockInfo(
-            name=name,
-            class_name=block.__class__.__name__,
-            description=getattr(block, "description", "") or "",
-            inputs=list(block.inputs) if hasattr(block, "inputs") else [],
-            outputs=list(block.intermediate_outputs) if hasattr(block, "intermediate_outputs") else [],
-        )
-        block_infos.append(info)
-    return block_infos
-
-
-def _get_type_name(type_hint):
-    if type_hint is None:
-        return None
-    if hasattr(type_hint, "__name__"):
-        return type_hint.__name__
-    if hasattr(type_hint, "__module__") and hasattr(type_hint, "__qualname__"):
-        return f"{type_hint.__module__}.{type_hint.__qualname__}"
-    return str(type_hint)
-
-
-def _is_internal_type(type_hint):
-    if type_hint is None:
-        return True
-    type_name = _get_type_name(type_hint)
-    if type_name is None:
-        return True
-    if type_name in INTERNAL_TYPE_NAMES or type_name in INTERNAL_TYPE_FULL_NAMES:
-        return True
-    type_str = str(type_hint)
-    for full_name in INTERNAL_TYPE_FULL_NAMES:
-        if full_name in type_str:
-            return True
-    if type_str.startswith("dict[") or type_str == "dict":
-        return True
-    return False
-
-
-def _type_hint_to_gradio(type_hint, param_name, default=None):
-    if _is_internal_type(type_hint):
-        return None
-
-    if param_name in SLIDER_PARAMS:
-        slider_opts = SLIDER_PARAMS[param_name]
-        val = default if default is not None else slider_opts.get("minimum", 0)
-        return (
-            f'gr.Slider(label="{param_name}", value={val!r}, '
-            f"minimum={slider_opts['minimum']}, maximum={slider_opts['maximum']}, "
-            f"step={slider_opts['step']})"
-        )
-
-    type_name = _get_type_name(type_hint)
-    type_str = str(type_hint)
-
-    if type_name == "str" or type_hint is str:
-        lines = 3 if "prompt" in param_name else 1
-        default_repr = f", value={default!r}" if default is not None else ""
-        return f'gr.Textbox(label="{param_name}", lines={lines}{default_repr})'
-
-    if type_name == "int" or type_hint is int:
-        val = f", value={default!r}" if default is not None else ""
-        return f'gr.Number(label="{param_name}", precision=0{val})'
-
-    if type_name == "float" or type_hint is float:
-        val = f", value={default!r}" if default is not None else ""
-        return f'gr.Number(label="{param_name}"{val})'
-
-    if type_name == "bool" or type_hint is bool:
-        val = default if default is not None else False
-        return f'gr.Checkbox(label="{param_name}", value={val!r})'
-
-    if "Image" in type_str:
-        if "list" in type_str.lower():
-            return f'gr.Gallery(label="{param_name}")'
-        return f'gr.Image(label="{param_name}")'
-
-    if default is not None:
-        return f'gr.Textbox(label="{param_name}", value={default!r})'
-
-    return f'gr.Textbox(label="{param_name}")'
-
-
-def _output_type_to_gradio(type_hint, param_name):
-    if _is_internal_type(type_hint):
-        return None
-    type_str = str(type_hint)
-    if "Image" in type_str:
-        if "list" in type_str.lower():
-            return f'gr.Gallery(label="{param_name}")'
-        return f'gr.Image(label="{param_name}")'
-    if type_hint is str:
-        return f'gr.Textbox(label="{param_name}")'
-    if type_hint is int or type_hint is float:
-        return f'gr.Number(label="{param_name}")'
-    return None
-
-
-def _classify_inputs(block_infos):
-    all_prior_outputs = {}
-
-    for info in block_infos:
-        user_inputs = []
-        port_connections = []
-        fixed_inputs = []
-
-        for inp in info.inputs:
-            if inp.name is None:
-                continue
-            if inp.name in all_prior_outputs:
-                port_connections.append((inp.name, all_prior_outputs[inp.name]))
-            elif _is_internal_type(inp.type_hint):
-                fixed_inputs.append(inp)
-            else:
-                user_inputs.append(inp)
-
-        info.user_inputs = user_inputs
-        info.port_connections = port_connections
-        info.fixed_inputs = fixed_inputs
-
-        for out in info.outputs:
-            if out.name and out.name not in all_prior_outputs:
-                all_prior_outputs[out.name] = info.name
-
-
-def _sanitize_name(name):
-    sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", name)
-    if sanitized and sanitized[0].isdigit():
-        sanitized = f"_{sanitized}"
-    return sanitized
-
-
-def _generate_code(block_infos, repo_id, blocks_class_name, workflow_label, workflow_resolve_code):
-    lines = []
-
-    lines.append(f'"""Daggr app for {blocks_class_name} ({workflow_label} workflow)')
-    lines.append("Generated by: diffusers-cli daggr")
-    lines.append('"""')
-    lines.append("")
-    lines.append("import gradio as gr")
-    lines.append("from daggr import FnNode, InputNode, Graph")
-    lines.append("")
-    lines.append("")
-
-    # Pipeline and resolved blocks loader
-    lines.append("_pipeline = None")
-    lines.append("_exec_blocks = None")
-    lines.append("")
-    lines.append("")
-    lines.append("def _get_pipeline():")
-    lines.append("    global _pipeline, _exec_blocks")
-    lines.append("    if _pipeline is None:")
-    lines.append("        from diffusers import ModularPipeline")
-    lines.append(f"        _pipeline = ModularPipeline.from_pretrained({repo_id!r}, trust_remote_code=True)")
-    lines.append("        _pipeline.load_components()")
-    lines.append(f"        _exec_blocks = {workflow_resolve_code}")
-    lines.append("    return _pipeline, _exec_blocks")
-    lines.append("")
-    lines.append("")
-
-    # Wrapper functions
-    for info in block_infos:
-        fn_name = f"run_{_sanitize_name(info.name)}"
-        all_input_names = []
-        for inp in info.inputs:
-            if inp.name is not None:
-                all_input_names.append(inp.name)
-
-        params = ", ".join(all_input_names)
-        lines.append(f"def {fn_name}({params}):")
-        lines.append("    from diffusers.modular_pipelines.modular_pipeline import PipelineState")
-        lines.append("")
-        lines.append("    pipe, exec_blocks = _get_pipeline()")
-        lines.append("    state = PipelineState()")
-        for inp_name in all_input_names:
-            lines.append(f'    state.set("{inp_name}", {inp_name})')
-        lines.append(f'    block = exec_blocks.sub_blocks["{info.name}"]')
-        lines.append("    _, state = block(pipe, state)")
-
-        if len(info.outputs) == 0:
-            lines.append("    return None")
-        elif len(info.outputs) == 1:
-            out = info.outputs[0]
-            lines.append(f'    return state.get("{out.name}")')
-        else:
-            out_names = [out.name for out in info.outputs]
-            out_dict = ", ".join(f'"{n}": state.get("{n}")' for n in out_names)
-            lines.append(f"    return {{{out_dict}}}")
-        lines.append("")
-        lines.append("")
-
-    # Collect all user-facing inputs across blocks
-    all_user_inputs = OrderedDict()
-    for info in block_infos:
-        for inp in info.user_inputs:
-            if inp.name not in all_user_inputs:
-                all_user_inputs[inp.name] = inp
-
-    # InputNode
-    if all_user_inputs:
-        lines.append("# -- User Inputs --")
-        lines.append('user_inputs = InputNode("User Inputs", ports={')
-        for inp_name, inp in all_user_inputs.items():
-            gradio_comp = _type_hint_to_gradio(inp.type_hint, inp_name, inp.default)
-            if gradio_comp:
-                lines.append(f'    "{inp_name}": {gradio_comp},')
-        lines.append("})")
-        lines.append("")
-        lines.append("")
-
-    # FnNode definitions
-    lines.append("# -- Pipeline Blocks --")
-    node_var_names = {}
-
-    for info in block_infos:
-        var_name = f"{_sanitize_name(info.name)}_node"
-        node_var_names[info.name] = var_name
-        fn_name = f"run_{_sanitize_name(info.name)}"
-
-        display_name = info.name.replace("_", " ").replace(".", " > ").title()
-
-        # Build inputs dict
-        input_entries = []
-        for inp in info.inputs:
-            if inp.name is None:
-                continue
-
-            connected = False
-            for conn_name, source_block in info.port_connections:
-                if conn_name == inp.name:
-                    source_var = node_var_names[source_block]
-                    input_entries.append(f'    "{inp.name}": {source_var}.{inp.name},')
-                    connected = True
-                    break
-
-            if not connected:
-                if inp.name in all_user_inputs:
-                    input_entries.append(f'    "{inp.name}": user_inputs.{inp.name},')
-                elif inp.default is not None:
-                    input_entries.append(f'    "{inp.name}": {inp.default!r},')
-                else:
-                    input_entries.append(f'    "{inp.name}": None,')
-
-        # Build outputs dict
-        output_entries = []
-        for out in info.outputs:
-            gradio_out = _output_type_to_gradio(out.type_hint, out.name)
-            if gradio_out:
-                output_entries.append(f'    "{out.name}": {gradio_out},')
-            else:
-                output_entries.append(f'    "{out.name}": None,')
-
-        lines.append(f"{var_name} = FnNode(")
-        lines.append(f"    fn={fn_name},")
-        lines.append(f'    name="{display_name}",')
-
-        if input_entries:
-            lines.append("    inputs={")
-            lines.extend(input_entries)
-            lines.append("    },")
-
-        if output_entries:
-            lines.append("    outputs={")
-            lines.extend(output_entries)
-            lines.append("    },")
-
-        lines.append(")")
-        lines.append("")
-
-    # Graph
-    lines.append("")
-    lines.append("# -- Graph --")
-    all_node_vars = []
-    if all_user_inputs:
-        all_node_vars.append("user_inputs")
-    all_node_vars.extend(node_var_names[info.name] for info in block_infos)
-
-    graph_name = f"{blocks_class_name} - {workflow_label}"
-    nodes_str = ", ".join(all_node_vars)
-    lines.append(f'graph = Graph("{graph_name}", nodes=[{nodes_str}])')
-    lines.append("graph.launch()")
-    lines.append("")
-
-    return "\n".join(lines)

src/diffusers/commands/diffusers_cli.py

@@ -16,7 +16,6 @@
 from argparse import ArgumentParser
 
 from .custom_blocks import CustomBlocksCommand
-from .daggr_app import DaggrCommand
 from .env import EnvironmentCommand
 from .fp16_safetensors import FP16SafetensorsCommand
 
@@ -29,7 +28,6 @@ def main():
     EnvironmentCommand.register_subcommand(commands_parser)
     FP16SafetensorsCommand.register_subcommand(commands_parser)
     CustomBlocksCommand.register_subcommand(commands_parser)
-    DaggrCommand.register_subcommand(commands_parser)
 
     # Let's go
     args = parser.parse_args()

src/diffusers/quantizers/quantization_config.py

@@ -36,7 +36,7 @@ from typing import Any, Callable
 
 from packaging import version
 
-from ..utils import deprecate, is_torch_available, is_torchao_available, is_torchao_version, logging
+from ..utils import is_torch_available, is_torchao_available, is_torchao_version, logging
 
 
 if is_torch_available():
@@ -844,8 +844,6 @@ class QuantoConfig(QuantizationConfigMixin):
         modules_to_not_convert: list[str] | None = None,
         **kwargs,
     ):
-        deprecation_message = "`QuantoConfig` is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoConfig", "1.0.0", deprecation_message)
         self.quant_method = QuantizationMethod.QUANTO
         self.weights_dtype = weights_dtype
         self.modules_to_not_convert = modules_to_not_convert

src/diffusers/quantizers/quanto/quanto_quantizer.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any
 from diffusers.utils.import_utils import is_optimum_quanto_version
 
 from ...utils import (
-    deprecate,
     get_module_from_name,
     is_accelerate_available,
     is_accelerate_version,
@@ -43,9 +42,6 @@ class QuantoQuantizer(DiffusersQuantizer):
         super().__init__(quantization_config, **kwargs)
 
     def validate_environment(self, *args, **kwargs):
-        deprecation_message = "The Quanto quantizer is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoQuantizer", "1.0.0", deprecation_message)
-
         if not is_optimum_quanto_available():
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"

tests/models/transformers/test_models_transformer_z_image.py

@@ -1,4 +1,3 @@
-# coding=utf-8
 # Copyright 2025 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,16 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import os
-import unittest
 
+import pytest
 import torch
 
 from diffusers import ZImageTransformer2DModel
+from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import IS_GITHUB_ACTIONS, torch_device
-from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
+from ...testing_utils import assert_tensors_close, torch_device
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    LoraTesterMixin,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TorchCompileTesterMixin,
+    TrainingTesterMixin,
+)
 
 
 # Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations
@@ -36,44 +42,38 @@ if hasattr(torch.backends, "cuda"):
     torch.backends.cuda.matmul.allow_tf32 = False
 
 
-@unittest.skipIf(
-    IS_GITHUB_ACTIONS,
-    reason="Skipping test-suite inside the CI because the model has `torch.empty()` inside of it during init and we don't have a clear way to override it in the modeling tests.",
-)
-class ZImageTransformerTests(ModelTesterMixin, unittest.TestCase):
-    model_class = ZImageTransformer2DModel
-    main_input_name = "x"
-    # We override the items here because the transformer under consideration is small.
-    model_split_percents = [0.9, 0.9, 0.9]
+def _concat_list_output(output):
+    """Model output `sample` is a list of tensors. Concatenate them for comparison."""
+    return torch.cat([t.flatten() for t in output])
 
-    def prepare_dummy_input(self, height=16, width=16):
-        batch_size = 1
-        num_channels = 16
-        embedding_dim = 16
-        sequence_length = 16
 
-        hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)]
-        encoder_hidden_states = [
-            torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size)
-        ]
-        timestep = torch.tensor([0.0]).to(torch_device)
-
-        return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep}
+class ZImageTransformerTesterConfig(BaseModelTesterConfig):
+    @property
+    def model_class(self):
+        return ZImageTransformer2DModel
 
     @property
-    def dummy_input(self):
-        return self.prepare_dummy_input()
-
-    @property
-    def input_shape(self):
+    def output_shape(self) -> tuple[int, ...]:
         return (4, 32, 32)
 
     @property
-    def output_shape(self):
+    def input_shape(self) -> tuple[int, ...]:
         return (4, 32, 32)
 
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = {
+    @property
+    def model_split_percents(self) -> list:
+        return [0.9, 0.9, 0.9]
+
+    @property
+    def main_input_name(self) -> str:
+        return "x"
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self):
+        return {
             "all_patch_size": (2,),
             "all_f_patch_size": (1,),
             "in_channels": 16,
@@ -89,83 +89,223 @@ class ZImageTransformerTests(ModelTesterMixin, unittest.TestCase):
             "axes_dims": [8, 4, 4],
             "axes_lens": [256, 32, 32],
         }
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
 
-    def setUp(self):
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
+    def get_dummy_inputs(self) -> dict[str, torch.Tensor | list]:
+        batch_size = 1
+        num_channels = 16
+        embedding_dim = 16
+        sequence_length = 16
+        height = 16
+        width = 16
+
+        hidden_states = [
+            randn_tensor((num_channels, 1, height, width), generator=self.generator, device=torch_device)
+            for _ in range(batch_size)
+        ]
+        encoder_hidden_states = [
+            randn_tensor((sequence_length, embedding_dim), generator=self.generator, device=torch_device)
+            for _ in range(batch_size)
+        ]
+        timestep = torch.tensor([0.0]).to(torch_device)
+
+        return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep}
+
+
+class TestZImageTransformer(ZImageTransformerTesterConfig, ModelTesterMixin):
+    """Core model tests for Z-Image Transformer."""
+
+    @torch.no_grad()
+    def test_determinism(self, atol=1e-5, rtol=0):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
+
+        inputs_dict = self.get_dummy_inputs()
+        first = _concat_list_output(model(**inputs_dict, return_dict=False)[0])
+        second = _concat_list_output(model(**inputs_dict, return_dict=False)[0])
+
+        mask = ~(torch.isnan(first) | torch.isnan(second))
+        assert_tensors_close(
+            first[mask], second[mask], atol=atol, rtol=rtol, msg="Model outputs are not deterministic"
+        )
+
+    def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5):
         torch.manual_seed(0)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(0)
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
+
+        model.save_pretrained(tmp_path)
+        new_model = self.model_class.from_pretrained(tmp_path)
+        new_model.to(torch_device)
+
+        for param_name in model.state_dict().keys():
+            param_1 = model.state_dict()[param_name]
+            param_2 = new_model.state_dict()[param_name]
+            assert param_1.shape == param_2.shape
+
+        inputs_dict = self.get_dummy_inputs()
+        image = _concat_list_output(model(**inputs_dict, return_dict=False)[0])
+        new_image = _concat_list_output(new_model(**inputs_dict, return_dict=False)[0])
+
+        assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.")
+
+    @torch.no_grad()
+    def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
+
+        model.save_pretrained(tmp_path, variant="fp16")
+        new_model = self.model_class.from_pretrained(tmp_path, variant="fp16")
+
+        with pytest.raises(OSError) as exc_info:
+            self.model_class.from_pretrained(tmp_path)
+
+        assert "Error no file named diffusion_pytorch_model.bin found in directory" in str(exc_info.value)
+
+        new_model.to(torch_device)
+
+        inputs_dict = self.get_dummy_inputs()
+        image = _concat_list_output(model(**inputs_dict, return_dict=False)[0])
+        new_image = _concat_list_output(new_model(**inputs_dict, return_dict=False)[0])
+
+        assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.")
+
+    @pytest.mark.skip("Model output `sample` is a list of tensors, not a single tensor.")
+    def test_outputs_equivalence(self, atol=1e-5, rtol=0):
+        pass
+
+    def test_sharded_checkpoints_with_parallel_loading(self, tmp_path, atol=1e-5, rtol=0):
+        from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME, constants
+
+        from ..testing_utils.common import calculate_expected_num_shards, compute_module_persistent_sizes
 
-    def tearDown(self):
-        super().tearDown()
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
         torch.manual_seed(0)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(0)
+        config = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+        model = self.model_class(**config).eval()
+        model = model.to(torch_device)
 
-    def test_gradient_checkpointing_is_applied(self):
-        expected_set = {"ZImageTransformer2DModel"}
-        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+        base_output = _concat_list_output(model(**inputs_dict, return_dict=False)[0])
 
-    @unittest.skip("Test is not supported for handling main inputs that are lists.")
-    def test_training(self):
-        super().test_training()
+        model_size = compute_module_persistent_sizes(model)[""]
+        max_shard_size = int((model_size * 0.75) / (2**10))
 
-    @unittest.skip("Test is not supported for handling main inputs that are lists.")
-    def test_ema_training(self):
-        super().test_ema_training()
+        original_parallel_loading = constants.HF_ENABLE_PARALLEL_LOADING
+        original_parallel_workers = getattr(constants, "HF_PARALLEL_WORKERS", None)
 
-    @unittest.skip("Test is not supported for handling main inputs that are lists.")
-    def test_effective_gradient_checkpointing(self):
-        super().test_effective_gradient_checkpointing()
+        try:
+            model.cpu().save_pretrained(tmp_path, max_shard_size=f"{max_shard_size}KB")
+            assert os.path.exists(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME))
 
-    @unittest.skip(
-        "Test needs to be revisited. But we need to ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices."
+            expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME))
+            actual_num_shards = len([file for file in os.listdir(tmp_path) if file.endswith(".safetensors")])
+            assert actual_num_shards == expected_num_shards
+
+            constants.HF_ENABLE_PARALLEL_LOADING = False
+            self.model_class.from_pretrained(tmp_path).eval().to(torch_device)
+
+            constants.HF_ENABLE_PARALLEL_LOADING = True
+            constants.DEFAULT_HF_PARALLEL_LOADING_WORKERS = 2
+
+            torch.manual_seed(0)
+            model_parallel = self.model_class.from_pretrained(tmp_path).eval()
+            model_parallel = model_parallel.to(torch_device)
+
+            output_parallel = _concat_list_output(model_parallel(**inputs_dict, return_dict=False)[0])
+
+            assert_tensors_close(
+                base_output, output_parallel, atol=atol, rtol=rtol, msg="Output should match with parallel loading"
+            )
+        finally:
+            constants.HF_ENABLE_PARALLEL_LOADING = original_parallel_loading
+            if original_parallel_workers is not None:
+                constants.HF_PARALLEL_WORKERS = original_parallel_workers
+
+
+class TestZImageTransformerMemory(ZImageTransformerTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests for Z-Image Transformer."""
+
+    @pytest.mark.skip(
+        "Ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices."
     )
     def test_layerwise_casting_training(self):
-        super().test_layerwise_casting_training()
-
-    @unittest.skip("Test is not supported for handling main inputs that are lists.")
-    def test_outputs_equivalence(self):
-        super().test_outputs_equivalence()
-
-    @unittest.skip("Test will pass if we change to deterministic values instead of empty in the DiT.")
-    def test_group_offloading(self):
-        super().test_group_offloading()
-
-    @unittest.skip("Test will pass if we change to deterministic values instead of empty in the DiT.")
-    def test_group_offloading_with_disk(self):
-        super().test_group_offloading_with_disk()
+        pass
 
 
-class ZImageTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
-    model_class = ZImageTransformer2DModel
-    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]
+class TestZImageTransformerTraining(ZImageTransformerTesterConfig, TrainingTesterMixin):
+    """Training tests for Z-Image Transformer."""
 
-    def prepare_init_args_and_inputs_for_common(self):
-        return ZImageTransformerTests().prepare_init_args_and_inputs_for_common()
+    def test_gradient_checkpointing_is_applied(self):
+        super().test_gradient_checkpointing_is_applied(expected_set={"ZImageTransformer2DModel"})
 
-    def prepare_dummy_input(self, height, width):
-        return ZImageTransformerTests().prepare_dummy_input(height=height, width=width)
+    @pytest.mark.skip("Test is not supported for handling main inputs that are lists.")
+    def test_training(self):
+        pass
 
-    @unittest.skip(
-        "The repeated block in this model is ZImageTransformerBlock, which is used for noise_refiner, context_refiner, and layers. As a consequence of this, the inputs recorded for the block would vary during compilation and full compilation with fullgraph=True would trigger recompilation at least thrice."
+    @pytest.mark.skip("Test is not supported for handling main inputs that are lists.")
+    def test_training_with_ema(self):
+        pass
+
+    @pytest.mark.skip("Test is not supported for handling main inputs that are lists.")
+    def test_gradient_checkpointing_equivalence(self, loss_tolerance=1e-5, param_grad_tol=5e-5, skip=None):
+        pass
+
+
+class TestZImageTransformerLoRA(ZImageTransformerTesterConfig, LoraTesterMixin):
+    """LoRA adapter tests for Z-Image Transformer."""
+
+    @pytest.mark.skip("Model output `sample` is a list of tensors, not a single tensor.")
+    def test_save_load_lora_adapter(self, tmp_path, rank=4, lora_alpha=4, use_dora=False, atol=1e-4, rtol=1e-4):
+        pass
+
+
+# TODO: Add pretrained_model_name_or_path once a tiny Z-Image model is available on the Hub
+# class TestZImageTransformerBitsAndBytes(ZImageTransformerTesterConfig, BitsAndBytesTesterMixin):
+#     """BitsAndBytes quantization tests for Z-Image Transformer."""
+
+
+# TODO: Add pretrained_model_name_or_path once a tiny Z-Image model is available on the Hub
+# class TestZImageTransformerTorchAo(ZImageTransformerTesterConfig, TorchAoTesterMixin):
+#     """TorchAo quantization tests for Z-Image Transformer."""
+
+
+class TestZImageTransformerCompile(ZImageTransformerTesterConfig, TorchCompileTesterMixin):
+    """Torch compile tests for Z-Image Transformer."""
+
+    @property
+    def different_shapes_for_compilation(self):
+        return [(4, 4), (4, 8), (8, 8)]
+
+    def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict[str, torch.Tensor | list]:
+        batch_size = 1
+        num_channels = 16
+        embedding_dim = 16
+        sequence_length = 16
+
+        hidden_states = [
+            randn_tensor((num_channels, 1, height, width), generator=self.generator, device=torch_device)
+            for _ in range(batch_size)
+        ]
+        encoder_hidden_states = [
+            randn_tensor((sequence_length, embedding_dim), generator=self.generator, device=torch_device)
+            for _ in range(batch_size)
+        ]
+        timestep = torch.tensor([0.0]).to(torch_device)
+
+        return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep}
+
+    @pytest.mark.skip(
+        "The repeated block in this model is ZImageTransformerBlock, which is used for noise_refiner, context_refiner, and layers. The inputs recorded for the block would vary during compilation and full compilation with fullgraph=True would trigger recompilation at least thrice."
     )
     def test_torch_compile_recompilation_and_graph_break(self):
-        super().test_torch_compile_recompilation_and_graph_break()
+        pass
 
-    @unittest.skip("Fullgraph AoT is broken")
-    def test_compile_works_with_aot(self):
-        super().test_compile_works_with_aot()
+    @pytest.mark.skip("Fullgraph AoT is broken")
+    def test_compile_works_with_aot(self, tmp_path):
+        pass
 
-    @unittest.skip("Fullgraph is broken")
+    @pytest.mark.skip("Fullgraph is broken")
     def test_compile_on_different_shapes(self):
-        super().test_compile_on_different_shapes()
+        pass