added testcase

Author: ishan-modi

src/diffusers/loaders/single_file_utils.py

@@ -182,7 +182,7 @@
     "hunyuan-video": {"pretrained_model_name_or_path": "hunyuanvideo-community/HunyuanVideo"},
     "instruct-pix2pix": {"pretrained_model_name_or_path": "timbrooks/instruct-pix2pix"},
     "lumina2": {"pretrained_model_name_or_path": "Alpha-VLLM/Lumina-Image-2.0"},
-    "sana":  {"pretrained_model_name_or_path": "Efficient-Large-Model/Sana_1600M_1024px"},
+    "sana":  {"pretrained_model_name_or_path": "Efficient-Large-Model/Sana_1600M_1024px_diffusers"},
 }
 
 # Use to configure model sample size when original config is provided
@@ -2878,6 +2878,7 @@ def convert_sana_transformer_to_diffusers(checkpoint, **kwargs):
 
     num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "blocks" in k))[-1] + 1  # noqa: C401
 
+
     # Positional and patch embeddings.
     checkpoint.pop("pos_embed")
     converted_state_dict["patch_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight")
@@ -2892,6 +2893,7 @@ def convert_sana_transformer_to_diffusers(checkpoint, **kwargs):
     converted_state_dict["time_embed.linear.bias"] = checkpoint.pop("t_block.1.bias")
 
     # Caption Projection.
+    checkpoint.pop("y_embedder.y_embedding")
     converted_state_dict["caption_proj.linear_1.weight"] = checkpoint.pop("y_embedder.y_proj.fc1.weight")
     converted_state_dict["caption_proj.linear_1.bias"] = checkpoint.pop("y_embedder.y_proj.fc1.bias")
     converted_state_dict["caption_proj.linear_2.weight"] = checkpoint.pop("y_embedder.y_proj.fc2.weight")

tests/single_file/test_sana_transformer.py

@@ -0,0 +1,37 @@
+import gc
+import unittest
+
+import torch
+
+from diffusers import (
+    SanaTransformer2DModel,
+)
+from diffusers.utils.testing_utils import (
+    backend_empty_cache,
+    enable_full_determinism,
+    require_torch_accelerator,
+    torch_device,
+)
+
+
+enable_full_determinism()
+
+
+@require_torch_accelerator
+class SanaTransformer2DModelSingleFileTests(unittest.TestCase):
+    model_class = SanaTransformer2DModel
+
+    repo_id = "Efficient-Large-Model/Sana_1600M_1024px_diffusers"
+
+    def setUp(self):
+        super().setUp()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def tearDown(self):
+        super().tearDown()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def test_single_file_components(self):
+        _ = self.model_class.from_pretrained(self.repo_id, subfolder="transformer")