Sana fails with BFloat16 and tiled VAE decode

[vladmandic]

Describe the bug

tried new sana 4k model, it fails when running in bfloat16 during vae decode if tiled decode is enabled.
issue is that in SanaMultiscaleAttnProcessor2_0 it only does upcasting conditionally:

        if use_linear_attention:
            # for linear attention upcast hidden_states to float32
            hidden_states = hidden_states.to(dtype=torch.float32)

and this remains true for entire execution when not using tiled decode
but when using tiled decode it may result in quadratic attention instead of linear and it ALSO needs upcasting.

Reproduction

run sana 4k model in bf16 with tiled vae.

Logs

│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/autoencoders/autoencoder_dc.py:93 in forward                                                                                                                                                                                                                                                                                                           │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│    92 │   def forward(self, x: torch.Tensor) -> torch.Tensor:                                                                                                                                                                                                                                                                                                                                                                    │
│ ❱  93 │   │   x = self.attn(x)                                                                                                                                                                                                                                                                                                                                                                                                   │
│    94 │   │   x = self.conv_out(x)                                                                                                                                                                                                                                                                                                                                                                                               │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736 in _wrapped_call_impl                                                                                                                                                                                                                                                                                                                   │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│   1735 │   │   else:                                                                                                                                                                                                                                                                                                                                                                                                             │
│ ❱ 1736 │   │   │   return self._call_impl(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                                                       │
│   1737                                                                                                                                                                                                                                                                                                                                                                                                                           │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747 in _call_impl                                                                                                                                                                                                                                                                                                                           │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│   1746 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                                                                                                                                                                                                                                                                                                                                                   │
│ ❱ 1747 │   │   │   return forward_call(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                                                          │
│   1748                                                                                                                                                                                                                                                                                                                                                                                                                           │
│                                                                                                                          
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/attention_processor.py:906 in forward                                                                                                                                                                                                                                                                                                                  │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│    905 │   def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:                                                                                                                                                                                                                                                                                                                                                       │
│ ❱  906 │   │   return self.processor(self, hidden_states)                                                                                                                                                                                                                                                                                                                                                                        │
│    907                                                                                                                                                                                                                                                                                                                                                                                                                           │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/attention_processor.py:5810 in __call__                                                                                                                                                                                                                                                                                                                │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│   5809 │   │   else:                                                                                                                                                                                                                                                                                                                                                                                                             │
│ ❱ 5810 │   │   │   hidden_states = attn.apply_quadratic_attention(query, key, value)                                                                                                                                                                                                                                                                                                                                             │
│   5811                                                                                                                                                                                                                                                                                                                                                                                                                           │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/attention_processor.py:902 in apply_quadratic_attention                                                                                                                                                                                                                                                                                                │
│                                                                                                                                                                                                                                                                                                                                                                                                                                  │
│    901 │   │   scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps)                                                                                                                                                                                                                                                                                                                                             │
│ ❱  902 │   │   hidden_states = torch.matmul(value, scores)                                                                                                                                                                                                                                                                                                                                                                       │
│    903 │   │   return hidden_states
RuntimeError: expected scalar type Float but found BFloat16

System Info

diffusers==c944f06
python==3.12.3

Who can help?

cc: @yiyixuxu @sayakpaul @DN6 @lawrence-cj

[nitinmukesh]

Same issue here. Build diffusers from source

Without this CUDA OOM, but inference works. pipe_sana.enable_vae_slicing() is working

None of these work
pipe_sana.vae.enable_tiling(tile_sample_min_height=1024, tile_sample_min_width=1024)

OR

pipe_sana.enable_vae_tiling()

  File "C:\aitest\diffuser_t2i_webui\venv\lib\site-packages\diffusers\models\attention_processor.py", line 902, in apply_quadratic_attention
    hidden_states = torch.matmul(value, scores)
RuntimeError: expected scalar type Float but found BFloat16

[a-r-r-o-w]

@vladmandic @nitinmukesh Could you try the linked PR above?

[lawrence-cj]

I have tested, the PR fixed the problem well. Please, feel free to merge it.
@a-r-r-o-w

[nitinmukesh]

@a-r-r-o-w @lawrence-cj

diffusers 0.33.0.dev0

pip install git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head

(venv) C:\aitest\diffuser_t2i_webui>pip install git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head
Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com
Collecting git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head
  Cloning https://github.com/huggingface/diffusers.git (to revision refs/pull/10595/head) to c:\users\nitin\appdata\local\temp\pip-req-build-917az2b1

  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done

Test 1: Sana 4K inference with enable_tiling

pipe_sana.vae.enable_tiling(tile_sample_min_height=1024, tile_sample_min_width=1024)

  File "C:\aitest\diffuser_t2i_webui\venv\lib\site-packages\diffusers\models\attention_processor.py", line 902, in apply_quadratic_attention
    hidden_states = torch.matmul(value, scores)
RuntimeError: expected scalar type Float but found BFloat16

Test 2: Sana 4K inference with enable_vae_tiling

pipe_sana.enable_vae_tiling()

line 902, in apply_quadratic_attention
    hidden_states = torch.matmul(value, scores)
RuntimeError: expected scalar type Float but found BFloat16

[lawrence-cj]

pipe.vae.enable_tiling(
    tile_sample_min_height=1024,
    tile_sample_min_width=1024,
    tile_sample_stride_height=896,
    tile_sample_stride_width=896,
)

try this one? @nitinmukesh

[nitinmukesh]

And here is the reproducible code

import torch
import numpy as np
import os
from datetime import datetime
from diffusers import SanaPipeline


pipeline_params = {
	"variant": "bf16",
	"torch_dtype": torch.bfloat16,
	"use_safetensors": True,
	"pretrained_model_name_or_path": "Efficient-Large-Model/Sana_1600M_4Kpx_BF16_diffusers"
}

# Initialize pipeline
pipe = SanaPipeline.from_pretrained(**pipeline_params)
pipe.to("cuda")
pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)
pipe.enable_model_cpu_offload()
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

image = pipe(
	prompt= "A cat",
	height= 4096,
	width= 4096,
	guidance_scale= 7.0,
	num_inference_steps= 30,
	generator= torch.Generator(device="cuda").manual_seed(0),
).images[0]
image.save("sana.png")

• pipe.enable_model_cpu_offload() - Works
• pipe.enable_vae_slicing() - Works
• pipe.enable_vae_tiling() - Doesn't work
• pipe_sana.vae.enable_tiling(tile_sample_min_height=1024, tile_sample_min_width=1024) - Doesn't work
• pipe.enable_sequential_cpu_offload() - Doesn't work
  RuntimeError: Tensor.item() cannot be called on meta tensors

[lawrence-cj]

And seems

pip install git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head

is not working.

I tried to use

pip install e .

under the diffusers dir. It works.

Also, I use the following to make it work:

pip install git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head#egg=diffusers

@nitinmukesh

[nitinmukesh]

@a-r-r-o-w @lawrence-cj

I can confirm this solution works (#10590 (comment))
Please merge the changes.

@lawrence-cj
This worked for me.
pip install git+https://github.com/huggingface/diffusers.git@refs/pull/10595/head

For whatever reason i just love this model. Thank you all for your help.

I never thought I would generate Full-HD images on 8 GB VRAM, and here we have 4K working. Big thanks to Diffusers and Sana team.

reduced quality for github

[a-r-r-o-w]

PR has been merged. @nitinmukesh Thanks for the kind words!