Upload matryoshka.py

matryoshka.py

@@ -3782,8 +3782,6 @@ class MatryoshkaPipeline(
         else:
             raise ValueError("Currently, nesting levels 0, 1, and 2 are supported.")
 
-        unet = unet.to(self.device)
-
         if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
             deprecation_message = (
                 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
@@ -3840,10 +3838,27 @@ class MatryoshkaPipeline(
             feature_extractor=feature_extractor,
             image_encoder=image_encoder,
         )
+        self.register_to_config(nesting_level=nesting_level)
         if hasattr(unet, "nest_ratio"):
             scheduler.scales = unet.nest_ratio + [1]
         self.image_processor = VaeImageProcessor(do_resize=False)
 
+    def change_nesting_level(self, nesting_level: int):
+        if nesting_level == 0:
+            self.unet = MatryoshkaUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                      subfolder="unet/nesting_level_0")
+            self.config.nesting_level = 0
+        elif nesting_level == 1:
+            self.unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                    subfolder="unet/nesting_level_1")
+            self.config.nesting_level = 1
+        elif nesting_level == 2:
+            self.unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                    subfolder="unet/nesting_level_2")
+            self.config.nesting_level = 2
+        else:
+            raise ValueError("Currently, nesting levels 0, 1, and 2 are supported.")
+
     def encode_prompt(
         self,
         prompt,