| from diffusers import DiffusionPipeline, LCMScheduler |
| import torch |
|
|
|
|
| def get_lcm_lora_pipeline( |
| base_model_id: str, |
| lcm_lora_id: str, |
| use_local_model: bool, |
| torch_data_type: torch.dtype, |
| ): |
| pipeline = DiffusionPipeline.from_pretrained( |
| base_model_id, |
| torch_dtype=torch_data_type, |
| local_files_only=use_local_model, |
| ) |
| pipeline.load_lora_weights( |
| lcm_lora_id, |
| local_files_only=use_local_model, |
| ) |
| if "lcm" in lcm_lora_id.lower(): |
| print("LCM LoRA model detected so using recommended LCMScheduler") |
| pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config) |
| pipeline.fuse_lora() |
| pipeline.unet.to(memory_format=torch.channels_last) |
| return pipeline |
|
|
