Skip to content

Deprecate upcast_vae in SDXL based pipelines #12619

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Deprecate upcast_vae in SDXL based pipelines #12619

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
21a03f9
update
DN6 Nov 10, 2025
7390638
update
DN6 Nov 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 2 additions & 14 deletions examples/community/lpw_stable_diffusion_xl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
TextualInversionLoaderMixin,
)
from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
from diffusers.models.attention_processor import AttnProcessor2_0, XFormersAttnProcessor
from diffusers.models.lora import adjust_lora_scale_text_encoder
from diffusers.pipelines.pipeline_utils import StableDiffusionMixin
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
Expand Down Expand Up @@ -1328,18 +1327,7 @@ def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, d

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

# Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
Expand Down Expand Up @@ -1921,7 +1909,7 @@ def denoising_value_valid(dnv):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

# unscale/denormalize the latents
Expand Down
2 changes: 1 addition & 1 deletion examples/community/masked_stable_diffusion_xl_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,7 +481,7 @@ def denoising_value_valid(dnv):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
Expand Down
2,455 changes: 1,218 additions & 1,237 deletions examples/community/mixture_tiling_sdxl.py
View file
Open in desktop

Large diffs are not rendered by default.

23 changes: 3 additions & 20 deletions examples/community/mod_controlnet_tile_sr_sdxl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,16 +39,13 @@
MultiControlNetModel,
UNet2DConditionModel,
)
from diffusers.models.attention_processor import (
AttnProcessor2_0,
XFormersAttnProcessor,
)
from diffusers.models.lora import adjust_lora_scale_text_encoder
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
from diffusers.schedulers import KarrasDiffusionSchedulers, LMSDiscreteScheduler
from diffusers.utils import (
USE_PEFT_BACKEND,
deprecate,
logging,
replace_example_docstring,
scale_lora_layers,
Expand Down Expand Up @@ -1222,21 +1219,7 @@ def prepare_tiles(

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

@property
def guidance_scale(self):
Expand Down Expand Up @@ -1820,7 +1803,7 @@ def __call__(
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

# unscale/denormalize the latents
Expand Down
22 changes: 2 additions & 20 deletions examples/community/pipeline_controlnet_xl_kolors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,6 @@
MultiControlNetModel,
UNet2DConditionModel,
)
from diffusers.models.attention_processor import (
AttnProcessor2_0,
XFormersAttnProcessor,
)
from diffusers.pipelines.kolors import ChatGLMModel, ChatGLMTokenizer
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
Expand Down Expand Up @@ -760,21 +756,7 @@ def _get_add_time_ids(

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

@property
def guidance_scale(self):
Expand Down Expand Up @@ -1331,7 +1313,7 @@ def _cn_patch_forward(*args, **kwargs):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

latents = latents / self.vae.config.scaling_factor
Expand Down
22 changes: 2 additions & 20 deletions examples/community/pipeline_controlnet_xl_kolors_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,6 @@
MultiControlNetModel,
UNet2DConditionModel,
)
from diffusers.models.attention_processor import (
AttnProcessor2_0,
XFormersAttnProcessor,
)
from diffusers.pipelines.kolors import ChatGLMModel, ChatGLMTokenizer
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
Expand Down Expand Up @@ -930,21 +926,7 @@ def _get_add_time_ids(

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

@property
def guidance_scale(self):
Expand Down Expand Up @@ -1533,7 +1515,7 @@ def _cn_patch_forward(*args, **kwargs):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

latents = latents / self.vae.config.scaling_factor
Expand Down
22 changes: 2 additions & 20 deletions examples/community/pipeline_controlnet_xl_kolors_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,6 @@
MultiControlNetModel,
UNet2DConditionModel,
)
from diffusers.models.attention_processor import (
AttnProcessor2_0,
XFormersAttnProcessor,
)
from diffusers.pipelines.kolors import ChatGLMModel, ChatGLMTokenizer
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
Expand Down Expand Up @@ -1006,21 +1002,7 @@ def _get_add_time_ids(

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

@property
def denoising_end(self):
Expand Down Expand Up @@ -1847,7 +1829,7 @@ def _cn_patch_forward(*args, **kwargs):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

latents = latents / self.vae.config.scaling_factor
Expand Down
19 changes: 4 additions & 15 deletions examples/community/pipeline_demofusion_sdxl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,11 @@
TextualInversionLoaderMixin,
)
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.models.attention_processor import AttnProcessor2_0, XFormersAttnProcessor
from diffusers.models.lora import adjust_lora_scale_text_encoder
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import (
deprecate,
is_accelerate_available,
is_accelerate_version,
is_invisible_watermark_available,
Expand Down Expand Up @@ -614,18 +614,7 @@ def tiled_decode(self, latents, current_height, current_width):

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
Expand Down Expand Up @@ -997,7 +986,7 @@ def __call__(
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
print("### Phase 1 Decoding ###")
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
Expand Down Expand Up @@ -1257,7 +1246,7 @@ def __call__(
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

print("### Phase {} Decoding ###".format(current_scale_num))
Expand Down
30 changes: 3 additions & 27 deletions examples/community/pipeline_faithdiff_stable_diffusion_xl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,6 @@
UNet2DConditionLoadersMixin,
)
from diffusers.models import AutoencoderKL
from diffusers.models.attention_processor import (
AttnProcessor2_0,
FusedAttnProcessor2_0,
LoRAAttnProcessor2_0,
LoRAXFormersAttnProcessor,
XFormersAttnProcessor,
)
from diffusers.models.lora import adjust_lora_scale_text_encoder
from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D, get_down_block
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
Expand Down Expand Up @@ -1637,24 +1630,7 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
return latents

def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.vae.decoder.mid_block.attentions[0].processor,
(
AttnProcessor2_0,
XFormersAttnProcessor,
LoRAXFormersAttnProcessor,
LoRAAttnProcessor2_0,
FusedAttnProcessor2_0,
),
)
# 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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

# Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
def get_guidance_scale_embedding(
Expand Down Expand Up @@ -1782,7 +1758,7 @@ def prepare_image_latents(
# needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
# if needs_upcasting:
# image = image.float()
# self.upcast_vae()
# self.vae.to(torch.float32)
self.unet.denoise_encoder.to(device=image.device, dtype=image.dtype)
image_latents = self.unet.denoise_encoder(image)
self.unet.denoise_encoder.to("cpu")
Expand Down Expand Up @@ -2235,7 +2211,7 @@ def __call__(
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
Expand Down
22 changes: 3 additions & 19 deletions examples/community/pipeline_kolors_differential_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,12 @@
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.loaders import IPAdapterMixin, StableDiffusionXLLoraLoaderMixin
from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
from diffusers.models.attention_processor import AttnProcessor2_0, FusedAttnProcessor2_0, XFormersAttnProcessor
from diffusers.pipelines.kolors.pipeline_output import KolorsPipelineOutput
from diffusers.pipelines.kolors.text_encoder import ChatGLMModel
from diffusers.pipelines.kolors.tokenizer import ChatGLMTokenizer
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor


Expand Down Expand Up @@ -711,22 +710,7 @@ def _get_add_time_ids(

# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.upcast_vae
def upcast_vae(self):
dtype = self.vae.dtype
self.vae.to(dtype=torch.float32)
use_torch_2_0_or_xformers = isinstance(
self.vae.decoder.mid_block.attentions[0].processor,
(
AttnProcessor2_0,
XFormersAttnProcessor,
FusedAttnProcessor2_0,
),
)
# 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:
self.vae.post_quant_conv.to(dtype)
self.vae.decoder.conv_in.to(dtype)
self.vae.decoder.mid_block.to(dtype)
deprecate("`upcast_vae` is deprecated")

# Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
def get_guidance_scale_embedding(
Expand Down Expand Up @@ -1252,7 +1236,7 @@ def denoising_value_valid(dnv):
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast

if needs_upcasting:
self.upcast_vae()
self.vae.to(torch.float32)
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
Expand Down
Loading
Loading