Unable to compile StableDiffusion
renderless opened this issue · comments
renderless commented
Describe the bug
Use following sample code to compile StableDiffusion model but compiler failed at export ot through torch.jit.script stage.
To Reproduce
- pull latest runtime image via
docker pull bladedisc/bladedisc:latest-runtime-torch1.13.1-cu116
- install packages as follows
diffusers==0.19.3
transformers==4.29.2
accelerate==0.21.0
- run example code
Example code
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="fp16",
torch_dtype=torch.float16).to("cuda")
import torch_blade
opt_cfg = torch_blade.Config()
opt_cfg.enable_fp16 = True
batch_size = 1
seq_len = 77
width = 512
height = 512
latent_channels = 4
ww = width // 8
hh = height // 8
# text encoder
encoder_ids = torch.ones(batch_size, seq_len, dtype=torch.int)
encoder_pos = torch.ones(batch_size, seq_len, dtype=torch.int)
encoder_inputs = (encoder_ids, encoder_pos)
# unet
unet_latent = torch.rand(batch_size, hh, ww, 4)
unet_timesteps = torch.rand(batch_size)
unet_text_embeddings = torch.rand(batch_size, seq_len, pipe.unet.config.cross_attention_dim)
unet_inputs = (unet_latent, unet_timesteps, unet_text_embeddings)
# vae decoder
decoder_inputs = torch.rand(batch_size, hh, ww, 4)
# todo: add encoder_inputs, unet_inputs, decoder_inputs
with opt_cfg, torch.no_grad():
encoder = torch_blade.optimize(pipe.text_encoder, model_inputs=encoder_inputs, allow_tracing=True)
unet = torch_blade.optimize(pipe.unet, model_inputs=unet_inputs, allow_tracing=True)
decoder = torch_blade.optimize(pipe.vae.decoder, model_inputs=decoder_inputs, allow_tracing=True)
Expected behavior
torch_blade.optimize should run without problems.