Upload matryoshka.py

scheduler/matryoshka.py

@@ -3746,10 +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__()
 
@@ -3801,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,
@@ -4510,7 +4522,7 @@ class MatryoshkaPipeline(
             timesteps, num_inference_steps = retrieve_timesteps(
                 self.scheduler, num_inference_steps, device, timesteps, sigmas
             )
-            timesteps = timesteps[:-1]  # is this correct???
+            timesteps = timesteps[:-1]
         else:
             timesteps = self.scheduler.timesteps