Upload matryoshka.py

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
@@ -3864,6 +3865,9 @@ class MatryoshkaPipeline(
         else:
             raise ValueError("Currently, nesting levels 0, 1, and 2 are supported.")
 
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def encode_prompt(
         self,
         prompt,
@@ -4645,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)[0]
+            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)