add VAE automatically tiling function;

Author: lawrence-cj

src/diffusers/pipelines/pag/pipeline_pag_sana.py

@@ -915,7 +915,13 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            try:
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            except torch.cuda.OutOfMemoryError as e:
+                print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
+                self.vae.enable_tiling(tile_sample_min_width=1024, tile_sample_min_height=1024)
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                self.vae.disable_tiling()
             if use_resolution_binning:
                 image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
 

src/diffusers/pipelines/sana/pipeline_sana.py

@@ -953,7 +953,13 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            try:
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            except torch.cuda.OutOfMemoryError as e:
+                print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
+                self.vae.enable_tiling(tile_sample_min_width=1024, tile_sample_min_height=1024)
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                self.vae.disable_tiling()
             if use_resolution_binning:
                 image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)