From e2bcbaeed0b5bfc93d0c7585db3b27b8b7581bf5 Mon Sep 17 00:00:00 2001
From: Disty0 <semihgulec2005@gmail.com>
Date: Fri, 10 Nov 2023 17:12:15 +0300
Subject: [PATCH] update lcm-convert.py

---
 CHANGELOG.md                           |  2 +-
 cli/{lcm_convert.py => lcm-convert.py} | 25 +++++++++++++++++++------
 modules/sd_samplers_diffusers.py       |  2 +-
 3 files changed, 21 insertions(+), 8 deletions(-)
 rename cli/{lcm_convert.py => lcm-convert.py} (62%)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f9f39f09..f9324bc7d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,7 +11,7 @@
     - set number of steps to some low number, for SD-XL 6-7 steps is normally sufficient  
       note: LCM scheduler does not support steps higher than 50
     - set cfg to 1 or 2  
-  - Add `cli/lcm_convert.py` script to convert any SD 1.5 or SD-XL model to LCM model  
+  - Add `cli/lcm-convert.py` script to convert any SD 1.5 or SD-XL model to LCM model  
     by baking in LORA and uploading to Huggingface, thanks @Disty0  
   - Add additional pipeline types for manual model loads when loading from `safetensors`  
   - Updated logic for calculating **steps** when using base/hires/refiner workflows  
diff --git a/cli/lcm_convert.py b/cli/lcm-convert.py
similarity index 62%
rename from cli/lcm_convert.py
rename to cli/lcm-convert.py
index 73cefa6fd..c2d7c266b 100644
--- a/cli/lcm_convert.py
+++ b/cli/lcm-convert.py
@@ -1,3 +1,4 @@
+import os
 import argparse
 import torch
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, AutoPipelineForText2Image, LCMScheduler
@@ -5,11 +6,13 @@ from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, AutoPi
 parser = argparse.ArgumentParser("lcm_convert")
 parser.add_argument("--name", help="Name of the new LCM model", type=str)
 parser.add_argument("--model", help="A model to convert", type=str)
+parser.add_argument("--lora-scale", default=1.0, help="Strenght of the LCM", type=float)
 parser.add_argument("--huggingface", action="store_true", help="Use Hugging Face models instead of safetensors models")
 parser.add_argument("--upload", action="store_true", help="Upload the new LCM model to Hugging Face")
-parser.add_argument("--no_save", action="store_true", help="Don't save the new LCM model to local disk")
+parser.add_argument("--no-half", action="store_true", help="Convert the new LCM model to FP32")
+parser.add_argument("--no-save", action="store_true", help="Don't save the new LCM model to local disk")
 parser.add_argument("--sdxl", action="store_true", help="Use SDXL models")
-parser.add_argument("--ssd_1b", action="store_true", help="Use SSD-1B models")
+parser.add_argument("--ssd-1b", action="store_true", help="Use SSD-1B models")
 
 args = parser.parse_args()
 
@@ -28,15 +31,25 @@ elif args.ssd_1b:
     pipeline.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
 else:
     pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
-pipeline.fuse_lora()
+pipeline.fuse_lora(lora_scale=args.lora_scale)
 
 #components = pipeline.components
 #pipeline = LatentConsistencyModelPipeline(**components)
 
-pipeline = pipeline.to(dtype=torch.float16)
+if args.no_half:
+    pipeline = pipeline.to(dtype=torch.float32)
+else:
+    pipeline = pipeline.to(dtype=torch.float16)
 print(pipeline)
 
 if not args.no_save:
-    pipeline.save_pretrained(args.name, variant="fp16")
+    os.makedirs(f"models--local--{args.name}/snapshots")
+    if args.no_half:
+        pipeline.save_pretrained(f"models--local--{args.name}/snapshots/{args.name}")
+    else:
+        pipeline.save_pretrained(f"models--local--{args.name}/snapshots/{args.name}", variant="fp16")
 if args.upload:
-    pipeline.push_to_hub(args.name, variant="fp16")
+    if args.no_half:
+        pipeline.push_to_hub(args.name)
+    else:
+        pipeline.push_to_hub(args.name, variant="fp16")
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index 6b1025d65..015a55aaa 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/modules/sd_samplers_diffusers.py
@@ -41,7 +41,7 @@ config = {
     'LMSD': { 'use_karras_sigmas': False, 'timestep_spacing': 'linspace', 'steps_offset': 0 },
     'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0 },
     'UniPC': { 'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'predict_x0': 'bh2', 'lower_order_final': True },
-    'LCM': { 'num_train_timesteps': 1000, 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False },      
+    'LCM': { 'num_train_timesteps': 1000, 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False },
 }
 
 samplers_data_diffusers = [