update

src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py

@@ -21,7 +21,6 @@ import torch
 from ...configuration_utils import FrozenDict
 from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKL
-from ...models.attention_processor import AttnProcessor2_0, XFormersAttnProcessor
 from ...utils import logging
 from ..modular_pipeline import (
     ModularPipelineBlocks,
@@ -74,25 +73,6 @@ class StableDiffusionXLDecodeStep(ModularPipelineBlocks):
             )
         ]
 
-    @staticmethod
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae with self->components
-    def upcast_vae(components):
-        dtype = components.vae.dtype
-        components.vae.to(dtype=torch.float32)
-        use_torch_2_0_or_xformers = isinstance(
-            components.vae.decoder.mid_block.attentions[0].processor,
-            (
-                AttnProcessor2_0,
-                XFormersAttnProcessor,
-            ),
-        )
-        # if xformers or torch_2_0 is used attention block does not need
-        # to be in float32 which can save lots of memory
-        if use_torch_2_0_or_xformers:
-            components.vae.post_quant_conv.to(dtype)
-            components.vae.decoder.conv_in.to(dtype)
-            components.vae.decoder.mid_block.to(dtype)
-
     @torch.no_grad()
     def __call__(self, components, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
@@ -103,7 +83,7 @@ class StableDiffusionXLDecodeStep(ModularPipelineBlocks):
             block_state.needs_upcasting = components.vae.dtype == torch.float16 and components.vae.config.force_upcast
 
             if block_state.needs_upcasting:
-                self.upcast_vae(components)
+                self.components.vae.to(torch.float32)
                 latents = latents.to(next(iter(components.vae.post_quant_conv.parameters())).dtype)
             elif latents.dtype != components.vae.dtype:
                 if torch.backends.mps.is_available():