Spaces:
Sleeping
Sleeping
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 | |