| 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] | |