Spaces:
Runtime error
Runtime error
File size: 933 Bytes
3e01790 6d3950c 3e01790 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
from config import MODEL_NAME,ADAPTER_NAME
import torch
from diffusers import DiffusionPipeline
from wandb.integration.diffusers import autolog
from config import PROJECT_NAME
autolog(init=dict(project=PROJECT_NAME))
def load_pipeline(model_name, adapter_name):
pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to(
"cuda"
)
pipe.load_lora_weights(adapter_name)
pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(
pipe.vae.decode, mode="max-autotune", fullgraph=True
)
pipe.fuse_qkv_projections()
return pipe
loaded_pipeline = load_pipeline(MODEL_NAME, ADAPTER_NAME)
images = loaded_pipeline('toaster', num_inference_steps=30).images[0]
|