|
import argparse |
|
import torch |
|
from library.custom_logging import setup_logging |
|
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, LCMScheduler |
|
from library.sdxl_model_util import convert_diffusers_unet_state_dict_to_sdxl, sdxl_original_unet, save_stable_diffusion_checkpoint, _load_state_dict_on_device as load_state_dict_on_device |
|
from accelerate import init_empty_weights |
|
|
|
|
|
logger = setup_logging() |
|
|
|
def parse_command_line_arguments(): |
|
argument_parser = argparse.ArgumentParser("lcm_convert") |
|
argument_parser.add_argument("--name", help="Name of the new LCM model", required=True, type=str) |
|
argument_parser.add_argument("--model", help="A model to convert", required=True, type=str) |
|
argument_parser.add_argument("--lora-scale", default=1.0, help="Strength of the LCM", type=float) |
|
argument_parser.add_argument("--sdxl", action="store_true", help="Use SDXL models") |
|
argument_parser.add_argument("--ssd-1b", action="store_true", help="Use SSD-1B models") |
|
return argument_parser.parse_args() |
|
|
|
def load_diffusion_pipeline(command_line_args): |
|
if command_line_args.sdxl or command_line_args.ssd_1b: |
|
return StableDiffusionXLPipeline.from_single_file(command_line_args.model) |
|
else: |
|
return StableDiffusionPipeline.from_single_file(command_line_args.model) |
|
|
|
def convert_and_save_diffusion_model(diffusion_pipeline, command_line_args): |
|
diffusion_pipeline.scheduler = LCMScheduler.from_config(diffusion_pipeline.scheduler.config) |
|
lora_weight_file_path = "latent-consistency/lcm-lora-" + ("sdxl" if command_line_args.sdxl else "ssd-1b" if command_line_args.ssd_1b else "sdv1-5") |
|
diffusion_pipeline.load_lora_weights(lora_weight_file_path) |
|
diffusion_pipeline.fuse_lora(lora_scale=command_line_args.lora_scale) |
|
|
|
diffusion_pipeline = diffusion_pipeline.to(dtype=torch.float16) |
|
logger.info("Saving file...") |
|
|
|
text_encoder_primary = diffusion_pipeline.text_encoder |
|
text_encoder_secondary = diffusion_pipeline.text_encoder_2 |
|
variational_autoencoder = diffusion_pipeline.vae |
|
unet_network = diffusion_pipeline.unet |
|
|
|
del diffusion_pipeline |
|
|
|
state_dict = convert_diffusers_unet_state_dict_to_sdxl(unet_network.state_dict()) |
|
with init_empty_weights(): |
|
unet_network = sdxl_original_unet.SdxlUNet2DConditionModel() |
|
|
|
load_state_dict_on_device(unet_network, state_dict, device="cuda", dtype=torch.float16) |
|
|
|
save_stable_diffusion_checkpoint( |
|
command_line_args.name, |
|
text_encoder_primary, |
|
text_encoder_secondary, |
|
unet_network, |
|
None, |
|
None, |
|
None, |
|
variational_autoencoder, |
|
None, |
|
None, |
|
torch.float16, |
|
) |
|
|
|
logger.info("...done saving") |
|
|
|
def main(): |
|
command_line_args = parse_command_line_arguments() |
|
try: |
|
diffusion_pipeline = load_diffusion_pipeline(command_line_args) |
|
convert_and_save_diffusion_model(diffusion_pipeline, command_line_args) |
|
except Exception as error: |
|
logger.error(f"An error occurred: {error}") |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
