Upload matryoshka.py

scheduler/matryoshka.py

@@ -20,6 +20,7 @@
 
 
 import inspect
+import gc
 import math
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -3753,7 +3754,7 @@ class MatryoshkaPipeline(
     """
 
     model_cpu_offload_seq = "text_encoder->image_encoder->unet"
-    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
+    _optional_components = ["unet", "safety_checker", "feature_extractor", "image_encoder"]
     _exclude_from_cpu_offload = ["safety_checker"]
     _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
 
@@ -3830,6 +3831,9 @@ class MatryoshkaPipeline(
             new_config["sample_size"] = 64
             unet._internal_dict = FrozenDict(new_config)
 
+        if hasattr(unet, "nest_ratio"):
+            scheduler.scales = unet.nest_ratio + [1]
+
         self.register_modules(
             text_encoder=text_encoder,
             tokenizer=tokenizer,
@@ -3838,10 +3842,32 @@ class MatryoshkaPipeline(
             feature_extractor=feature_extractor,
             image_encoder=image_encoder,
         )
-        if hasattr(unet, "nest_ratio"):
-            scheduler.scales = unet.nest_ratio + [1]
+        self.register_to_config(nesting_level=nesting_level)
         self.image_processor = VaeImageProcessor(do_resize=False)
 
+    def change_nesting_level(self, nesting_level: int):
+        if nesting_level == 0:
+            if hasattr(self.unet, "nest_ratio"):
+                self.scheduler.scales = None
+            self.unet = MatryoshkaUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                      subfolder="unet/nesting_level_0").to(self.device)
+            self.config.nesting_level = 0
+        elif nesting_level == 1:
+            self.unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                    subfolder="unet/nesting_level_1").to(self.device)
+            self.config.nesting_level = 1
+            self.scheduler.scales = self.unet.nest_ratio + [1]
+        elif nesting_level == 2:
+            self.unet = NestedUNet2DConditionModel.from_pretrained("tolgacangoz/matryoshka-diffusion-models",
+                                                                    subfolder="unet/nesting_level_2").to(self.device)
+            self.config.nesting_level = 2
+            self.scheduler.scales = self.unet.nest_ratio + [1]
+        else:
+            raise ValueError("Currently, nesting levels 0, 1, and 2 are supported.")
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def encode_prompt(
         self,
         prompt,
@@ -4623,9 +4649,12 @@ class MatryoshkaPipeline(
         image = latents
 
         if self.scheduler.scales is not None:
-            for i in range(len(image)):
-                image[i] = image[i] * self.scheduler.scales[i]
-                image[i] = self.image_processor.postprocess(image[i], output_type=output_type)
+            scales = [
+                image[i].size(-1) / image[-1].size(-1)
+                for i in range(len(image))
+                ]
+            for i, (img, scale) in enumerate(zip(image, scales)):
+                image[i] = self.image_processor.postprocess(img * scale, output_type=output_type)[0]
         else:
             image = self.image_processor.postprocess(image, output_type=output_type)