DCAE-Sana obtains the NAN value after decoding the latent which comes from encoding image

[ZouYa99]

I ran the following code.However, the "output" is all NAN values. I'm confused. Can you help me?

from efficientvit.ae_model_zoo import DCAE_HF
import torch
dc_ae = DCAE_HF.from_pretrained("/data/dc-ae-f32c32-sana-1.0")
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.utils import save_image
from efficientvit.apps.utils.image import DMCrop

device = torch.device("cuda")
dc_ae = dc_ae.to(device).eval()

transform = transforms.Compose([
    DMCrop(512), # resolution
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
image = Image.open("assets/fig/girl.png")
x = transform(image)[None].to(device)
latent = dc_ae.encode(x)
output = dc_ae.decode(latent)

[lawrence-cj]

Hi, can you check If the dc_ae.dtype is FP32 or BF16? @ZouYa99

[ZouYa99]

It is torch.float32. @lawrence-cj

[lawrence-cj]

@chenjy2003 Junyu, can you help here?

[chenjy2003]

Hi @ZouYa99 , I ran the same code (except replacing /data/dc-ae-f32c32-sana-1.0 with mit-han-lab/dc-ae-f32c32-sana-1.0), and the result is normal. I'm not sure if there is something wrong with this local checkpoint. Could you please elaborate on how you downloaded the weights and check if the weights are correct?

[ZouYa99]

I am sorry to say that my device cannot run "mit-han-lab/dc-ae-f32c32-sana-1.0" directly because of the network. So I downloaded this model again (https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0). But there's still a problem with NAN.

[chenjy2003]

@ZouYa99 I'm still not sure the reason for this NAN issue. Maybe something is wrong with the environment. Since our models are merged into diffusers, I would recommend you give it a try. Please upgrade diffusers first by pip install -U diffusers.

from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.utils import save_image
from diffusers import AutoencoderDC

device = torch.device("cuda")
dc_ae: AutoencoderDC = AutoencoderDC.from_pretrained(f"mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers", torch_dtype=torch.float32).to(device).eval()

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5),
])

image = Image.open("assets/fig/girl.png")
x = transform(image)[None].to(device)
latent = dc_ae.encode(x).latent
y = dc_ae.decode(latent).sample
save_image(y * 0.5 + 0.5, "demo_dc_ae.png")