lcm scheduler

modules/async_worker.py

@@ -575,6 +575,23 @@ def worker():
         preparation_time = time.perf_counter() - execution_start_time
         print(f'Preparation time: {preparation_time:.2f} seconds')
 
+        final_sampler_name = sampler_name
+        final_scheduler_name = scheduler_name
+
+        if scheduler_name == 'lcm':
+            final_scheduler_name = 'sgm_uniform'
+            if pipeline.final_unet is not None:
+                pipeline.final_unet = core.opModelSamplingDiscrete.patch(
+                    pipeline.final_unet,
+                    sampling='lcm',
+                    zsnr=False)[0]
+            if pipeline.final_refiner_unet is not None:
+                pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch(
+                    pipeline.final_refiner_unet,
+                    sampling='lcm',
+                    zsnr=False)[0]
+            print('Using lcm scheduler.')
+
         outputs.append(['preview', (13, 'Moving model to GPU ...', None)])
 
         def callback(step, x0, x, total_steps, y):
@@ -609,8 +626,8 @@ def worker():
                     height=height,
                     image_seed=task['task_seed'],
                     callback=callback,
-                    sampler_name=sampler_name,
-                    scheduler_name=scheduler_name,
+                    sampler_name=final_sampler_name,
+                    scheduler_name=final_scheduler_name,
                     latent=initial_latent,
                     denoise=denoising_strength,
                     tiled=tiled,

modules/core.py

@@ -26,6 +26,7 @@ from modules.patch import patched_sampler_cfg_function
 from fcbh.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
 from modules.config import path_embeddings
 from modules.lora import load_dangerous_lora
+from fcbh_extras.nodes_model_advanced import ModelSamplingDiscrete
 
 
 opEmptyLatentImage = EmptyLatentImage()
@@ -35,6 +36,7 @@ opVAEDecodeTiled = VAEDecodeTiled()
 opVAEEncodeTiled = VAEEncodeTiled()
 opControlNetApplyAdvanced = ControlNetApplyAdvanced()
 opFreeU = FreeU_V2()
+opModelSamplingDiscrete = ModelSamplingDiscrete()
 
 
 class StableDiffusionModel:

modules/flags.py

@@ -14,7 +14,7 @@ KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral"
                   "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
                   "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
 
-SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm"]
 SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
 
 sampler_list = SAMPLER_NAMES