diff --git a/tests/diffusion/models/magi2/test_native_preview.py b/tests/diffusion/models/magi2/test_native_preview.py index 9167d72ee..10458c493 100644 --- a/tests/diffusion/models/magi2/test_native_preview.py +++ b/tests/diffusion/models/magi2/test_native_preview.py @@ -90,6 +90,21 @@ def test_hsdp_policy_shards_individual_preview_layers(): assert not condition("pre_adapter", model.pre_adapter) +def test_layerwise_offload_contract() -> None: + from vllm_omni.diffusion.offloader import get_blocks_from_dit + + model = Magi2PreviewTransformer(_tiny_config(num_layers=2)) + + assert model._layerwise_offload_blocks_attrs == ["block"] + assert "block" in dict(model.named_children()) + assert list(model.block) == list(model.block.layers) + assert model.layers is model.block.layers + + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == ["block"] + assert blocks == list(model.block.layers) + + def test_data_proxy_keeps_output_layout_request_scoped() -> None: proxy = Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=0)) diff --git a/vllm_omni/diffusion/models/magi2/modeling_magi2.py b/vllm_omni/diffusion/models/magi2/modeling_magi2.py index 44e282aa0..0009a2ecc 100644 --- a/vllm_omni/diffusion/models/magi2/modeling_magi2.py +++ b/vllm_omni/diffusion/models/magi2/modeling_magi2.py @@ -670,11 +670,9 @@ def _is_magi2_transformer_layer(_name: str, module: nn.Module) -> bool: class Magi2PreviewTransformer(nn.Module): """Native preview DiT with the released checkpoint hierarchy.""" - # ``block`` must remain the registered child name because it is part of the - # released checkpoint hierarchy. The ``layers`` property below exposes its - # ModuleList through the standard layerwise-offload contract without - # changing the registered module names. - _layerwise_offload_blocks_attrs = ["layers"] + # ``block`` is the registered child module containing the repeated + # layers, exposed as the streamable blocks container for layerwise offload. + _layerwise_offload_blocks_attrs = ["block"] _hsdp_shard_conditions = [_is_magi2_transformer_layer] _hsdp_preserve_parameter_dtypes = True _EP_SHARDED_SUFFIXES = ( @@ -704,7 +702,7 @@ class Magi2PreviewTransformer(nn.Module): @property def layers(self) -> nn.ModuleList: - """Expose Preview layers to the shared offload block discovery API.""" + """Expose Preview layers for direct access and compatibility.""" return self.block.layers