Fix DreamLite legacy block type aliases (huggingface#14066)

* Fix DreamLite legacy block type aliases

* Cover DreamLite non-V1 up block alias

---------

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

Author: ElectricGoal

src/diffusers/models/unets/unet_dreamlite.py

@@ -1340,10 +1340,27 @@ def forward(
 # Local block dispatch (DreamLite-only)
 #
 # The string ``down_block_type`` / ``up_block_type`` / ``mid_block_type`` keys
-# persisted in saved checkpoints' ``config.json`` mirror the Python class names
-# defined above. The ``carlofkl/DreamLite-{base,mobile}`` Hub repos
-# (``diffusers`` branch) ship configs that use these exact keys.
+# persisted in saved checkpoints' ``config.json`` usually mirror the Python class
+# names defined above. Some configs use upstream UNet block names instead.
 # ---------------------------------------------------------------------------
+_DREAMLITE_DOWN_BLOCK_ALIASES = {
+    "CrossAttnDownRemoveSelfAttnBlock2D": "DreamLiteCrossAttnNoSelfAttnDownBlock2D",
+    "CrossAttnDownBlock2D": "DreamLiteCrossAttnDownBlock2D",
+    "DownBlock2D": "DreamLiteDownBlock2D",
+}
+
+_DREAMLITE_MID_BLOCK_ALIASES = {
+    "UNetMidBlock2DCrossAttn": "DreamLiteUNetMidBlock2DCrossAttn",
+}
+
+_DREAMLITE_UP_BLOCK_ALIASES = {
+    "CrossAttnUpRemoveSelfAttnBlock2D": "DreamLiteCrossAttnNoSelfAttnUpBlock2D",
+    "CrossAttnUpRemoveSelfAttnBlock2DV1": "DreamLiteCrossAttnNoSelfAttnUpBlock2D",
+    "CrossAttnUpBlock2D": "DreamLiteCrossAttnUpBlock2D",
+    "UpBlock2D": "DreamLiteUpBlock2D",
+}
+
+
 def _get_down_block_dreamlite(
     down_block_type: str,
     *,
@@ -1371,6 +1388,8 @@ def _get_down_block_dreamlite(
     ff_mult,
     num_kv_heads,
 ):
+    down_block_type = _DREAMLITE_DOWN_BLOCK_ALIASES.get(down_block_type, down_block_type)
+
     if down_block_type == "DreamLiteDownBlock2D":
         return DreamLiteDownBlock2D(
             num_layers=num_layers,
@@ -1447,6 +1466,8 @@ def _get_mid_block_dreamlite(
 ):
     if mid_block_type is None:
         return None
+    mid_block_type = _DREAMLITE_MID_BLOCK_ALIASES.get(mid_block_type, mid_block_type)
+
     if mid_block_type == "DreamLiteUNetMidBlock2DCrossAttn":
         return DreamLiteUNetMidBlock2DCrossAttn(
             transformer_layers_per_block=transformer_layers_per_block,
@@ -1501,6 +1522,8 @@ def _get_up_block_dreamlite(
     ff_mult,
     num_kv_heads,
 ):
+    up_block_type = _DREAMLITE_UP_BLOCK_ALIASES.get(up_block_type, up_block_type)
+
     if up_block_type == "DreamLiteUpBlock2D":
         return DreamLiteUpBlock2D(
             num_layers=num_layers,

tests/pipelines/dreamlite/test_pipeline_dreamlite.py

@@ -131,6 +131,79 @@ class DreamLitePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
     test_layerwise_casting = False
     test_group_offloading = False
 
+    def test_legacy_block_type_aliases(self):
+        unet = DreamLiteUNetModel(
+            sample_size=8,
+            in_channels=4,
+            out_channels=4,
+            down_block_types=(
+                "CrossAttnDownRemoveSelfAttnBlock2D",
+                "CrossAttnDownRemoveSelfAttnBlock2D",
+                "CrossAttnDownBlock2D",
+            ),
+            mid_block_type="UNetMidBlock2DCrossAttn",
+            up_block_types=(
+                "CrossAttnUpBlock2D",
+                "CrossAttnUpRemoveSelfAttnBlock2DV1",
+                "UpBlock2D",
+            ),
+            block_out_channels=(16, 32, 64),
+            cross_attention_dim=_CROSS_ATTN_DIM,
+            attention_head_dim=8,
+            layers_per_block=1,
+            norm_num_groups=8,
+            transformer_layers_per_block=1,
+        )
+
+        self.assertEqual(
+            [block.__class__.__name__ for block in unet.down_blocks],
+            [
+                "DreamLiteCrossAttnNoSelfAttnDownBlock2D",
+                "DreamLiteCrossAttnNoSelfAttnDownBlock2D",
+                "DreamLiteCrossAttnDownBlock2D",
+            ],
+        )
+        self.assertEqual(unet.mid_block.__class__.__name__, "DreamLiteUNetMidBlock2DCrossAttn")
+        self.assertEqual(
+            [block.__class__.__name__ for block in unet.up_blocks],
+            [
+                "DreamLiteCrossAttnUpBlock2D",
+                "DreamLiteCrossAttnNoSelfAttnUpBlock2D",
+                "DreamLiteUpBlock2D",
+            ],
+        )
+
+        unet_with_non_v1_up_alias = DreamLiteUNetModel(
+            sample_size=8,
+            in_channels=4,
+            out_channels=4,
+            down_block_types=(
+                "CrossAttnDownRemoveSelfAttnBlock2D",
+                "CrossAttnDownRemoveSelfAttnBlock2D",
+                "CrossAttnDownBlock2D",
+            ),
+            mid_block_type="UNetMidBlock2DCrossAttn",
+            up_block_types=(
+                "CrossAttnUpBlock2D",
+                "CrossAttnUpRemoveSelfAttnBlock2D",
+                "UpBlock2D",
+            ),
+            block_out_channels=(16, 32, 64),
+            cross_attention_dim=_CROSS_ATTN_DIM,
+            attention_head_dim=8,
+            layers_per_block=1,
+            norm_num_groups=8,
+            transformer_layers_per_block=1,
+        )
+        self.assertEqual(
+            [block.__class__.__name__ for block in unet_with_non_v1_up_alias.up_blocks],
+            [
+                "DreamLiteCrossAttnUpBlock2D",
+                "DreamLiteCrossAttnNoSelfAttnUpBlock2D",
+                "DreamLiteUpBlock2D",
+            ],
+        )
+
     def get_dummy_components(self):
         torch.manual_seed(0)
         unet = DreamLiteUNetModel(