Upload matryoshka.py

matryoshka.py

@@ -3746,11 +3746,12 @@ class MatryoshkaPipeline(
         self,
         text_encoder: T5EncoderModel,
         tokenizer: T5TokenizerFast,
-        unet: MatryoshkaUNet2DConditionModel,
         scheduler: MatryoshkaDDIMScheduler,
+        unet: MatryoshkaUNet2DConditionModel = None,
         feature_extractor: CLIPImageProcessor = None,
         image_encoder: CLIPVisionModelWithProjection = None,
         trust_remote_code: bool = False,
+        nesting_level: int = 0,
     ):
         super().__init__()
 
@@ -3802,6 +3803,16 @@ class MatryoshkaPipeline(
             new_config["sample_size"] = 64
             unet._internal_dict = FrozenDict(new_config)
 
+        if nesting_level == 0:
+            unet = MatryoshkaUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                  subfolder="unet/nesting_level_0")
+        elif nesting_level == 1:
+            unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                              subfolder="unet/nesting_level_1")
+        elif nesting_level == 2:
+            unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                              subfolder="unet/nesting_level_2")
+
         self.register_modules(
             text_encoder=text_encoder,
             tokenizer=tokenizer,