Input type (c10::Half) and bias type (float) mismatch in pipeline_stable_diffusion_xl.py

[Yukun-Huang]

Describe the bug

I'm using the official example and trying to call StableDiffusionXLPipeline (fp16) twice. But in the second call, an error message appeared:

RuntimeError: Input type (c10::Half) and bias type (float) should be the same.

There are similar issues like #4619 and #4478 but don't seem to apply in my case.

I found the reason for this error is from the following code:

diffusers/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

Lines 839 to 848 in 24c5e77

	# make sure the VAE is in float32 mode, as it overflows in float16
	if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
	self.upcast_vae()
	latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
	if not output_type == "latent":
	image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
	else:
	image = latents
	return StableDiffusionXLPipelineOutput(images=image)

For the second call of pipe, self.vae.dtype is already float32, so the type conversion in L840-L842 is skipped, which makes latents still be float16 and conflicts with float32 vae.

Reproduction

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16,
                                         use_safetensors=True, variant="fp16")
pipe.to("cuda")

prompt = "An astronaut riding a green horse"

pipe(prompt=prompt).images[0].save('1.jpg')

pipe(prompt=prompt).images[0].save('2.jpg')
# RuntimeError: Input type (c10::Half) and bias type (float) should be the same

Logs

No response

System Info

• diffusers version: 0.21.0.dev0
• Platform: Linux-5.4.0-122-generic-x86_64-with-glibc2.17
• Python version: 3.8.15
• PyTorch version (GPU?): 1.13.1 (True)
• Huggingface_hub version: 0.16.2
• Transformers version: 4.30.2
• Accelerate version: 0.18.0
• xFormers version: not installed
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

@williamberman, @patrickvonplaten, and @sayakpaul

[sayakpaul]

Unable to reproduce: https://colab.research.google.com/gist/sayakpaul/dd70999a6ef23ef57ad4a5672b5d038d/scratchpad.ipynb

[EnricoBeltramo]

For me I fixed using a fp16 VAE, i.e:
https://huggingface.co/madebyollin/sdxl-vae-fp16-fix

[perrywu1989]

same issue

[sayakpaul]

I need a reproducible Colab Notebook to confirm that.

As you can check my comment above, I am unable to reproduce the problem. So, I am going to close this issue. Feel free to reopen if you can reproduce it and provide a Colab Notebook :-)

[jamie0725]

@sayakpaul I can also reproduce the issue with the following codes:

import gradio as gr
from diffusers import StableDiffusionPipeline
from diffusers import DiffusionPipeline
import torch

lora_model_id = "saved_model"
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe = pipe.to("cuda")
# pipe.load_lora_weights(lora_model_id)

def generate(prompt):
    image = pipe(prompt, num_inference_steps=25, num_images_per_prompt=1).images
    
    return image[0]

with gr.Blocks() as demo:
    with gr.Row():
        text_input = gr.Textbox(label="Input Prompt", lines=2, placeholder="A photo of sks dog in the pot.")
        img_output = gr.Image()
    button = gr.Button("Generate")
    button.click(generate, inputs=text_input, outputs=img_output)

demo.launch(share=True)

[jamie0725]

a screenshot of the error:

[sayakpaul]

Open a new issue then without any gradio related component. Ensure we can reproduce the issue on a Colab Notebook kas this only concerns inference.

[Yukun-Huang]

@sayakpaul Thanks for the response. The issue can be reproduced with a torch version of 1.13.1:
https://colab.research.google.com/gist/hyk1996/82ab6af1a9ad3d0b82b1f599f877c96a/scratchpad.ipynb

[sayakpaul]

Interesting. If it's a Torch version problem, I wonder if it could be more minimally reproduced. That way we could directly report that on PyTorch GitHub no?

[Yukun-Huang]

I think it is an implementation bug in pipeline_stable_diffusion_xl.py as well as other xl-based pipelines. It forces vae in float32 mode, but latents in float16 are not converted to float32 in the second inference, as I described at the beginning.

The case of torch>=2.0 works because some layers of vae are not converted to float32 when using torch>=2.0 as belows:

diffusers/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

Lines 550 to 553 in b1290d3

	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)

so latents can be passed forward correctly, and then luckily encounter a type conversion in:

diffusers/src/diffusers/models/vae.py

Line 267 in b1290d3

sample = sample.to(upscale_dtype)

I made an ugly fix to this issue:

# L839-L842 in pipeline_stable_diffusion_xl.py [029fb41]
# make sure the VAE is in float32 mode, as it overflows in float16
if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
    self.upcast_vae()
    # latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)  # before
latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)        # after

[sayakpaul]

Do you maybe want to open a PR for us to discuss this further?

[Yukun-Huang]

Sure, I will open a PR later.