From 96a138fe3cfb64a07c4b9d1dc6cf0cf34c611298 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 31 Oct 2023 11:42:32 -0400 Subject: [PATCH] handle k-diffusion module --- README.md | 17 ++++++------ modules/sd_samplers_kdiffusion.py | 43 +++++++++++++++++++++---------- modules/shared.py | 2 +- 3 files changed, 40 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 8c120b481..84e66df8d 100644 --- a/README.md +++ b/README.md @@ -19,14 +19,14 @@ All Individual features are not listed here, instead check [ChangeLog](CHANGELOG.md) for full list of changes. -- Optimized processing with latest **torch** developments - Including built-in support for `torch.compile` - Support for multiple backends! **original** and **diffusers** - Support for multiple diffusion models! - Stable Diffusion, SD-XL, Kandinsky, DeepFloyd IF, UniDiffusion, SD-Distilled, etc. + **Stable Diffusion, SD-XL, LCM, Segmind, Kandinsky, Wuerstchen, DeepFloyd IF, UniDiffusion, SD-Distilled, etc.** - Fully multiplatform with platform specific autodetection and tuning performed on install - Windows / Linux / MacOS with CPU / nVidia / AMD / Intel / DirectML / OpenVINO + **Windows / Linux / MacOS with CPU / nVidia / AMD / IntelArc / DirectML / OpenVINO / ONNX+Olive** +- Optimized processing with latest **torch** developments + Including built-in support for `torch.compile` - Improved prompt parser - Enhanced *Lora*/*Locon*/*Lyco* code supporting latest trends in training - Built-in queue management @@ -35,7 +35,7 @@ All Individual features are not listed here, instead check [ChangeLog](CHANGELOG - Modern localization and hints engine - Broad compatibility with existing extensions ecosystem and new extensions manager - Built in installer with automatic updates and dependency management -- Modernized UI (still based on Gradio) with theme support +- Modernized UI with theme support and number of built-in themes ## Backend support @@ -55,12 +55,13 @@ Additional models will be added as they become available and there is public int - [Stable Diffusion](https://github.com/Stability-AI/stablediffusion/) 1.x and 2.x *(all variants)* - [Stable Diffusion XL](https://github.com/Stability-AI/generative-models) -- [Kandinsky](https://github.com/ai-forever/Kandinsky-2) 2.1 and 2.2 -- [SD-Distilled](https://huggingface.co/blog/sd_distillation) *(all variants)* -- [Wuerstchen](https://huggingface.co/blog/wuertschen) - [Segmind SSD-1B](https://huggingface.co/segmind/SSD-1B) +- [LCM: Latent Consistency Models](https://github.com/openai/consistency_models) +- [Kandinsky](https://github.com/ai-forever/Kandinsky-2) 2.1 and 2.2 +- [Wuerstchen](https://huggingface.co/blog/wuertschen) - [UniDiffusion](https://github.com/thu-ml/unidiffuser) - [DeepFloyd IF](https://github.com/deep-floyd/IF) +- [SD-Distilled](https://huggingface.co/blog/sd_distillation) *(all variants)* ## Platform support diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 1d56513b2..578953740 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -1,17 +1,36 @@ -from collections import deque +import sys +import time import inspect +from collections import deque import torch -import k_diffusion.sampling from modules import prompt_parser from modules import devices from modules import sd_samplers_common - import modules.shared as shared from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback from modules.script_callbacks import CFGDenoisedParams, cfg_denoised_callback from modules.script_callbacks import AfterCFGCallbackParams, cfg_after_cfg_callback +# deal with k-diffusion imports +k_sampling = None +try: + import k_diffusion.sampling as k_sampling # pylint: disable=wrong-import-order +except ImportError: + pass +try: + if k_sampling is None: + import importlib + k_diffusion = importlib.import_module('modules.k-diffusion.k_diffusion') + k_sampling = k_diffusion.sampling +except: + pass +if k_sampling is None: + shared.log.info(f'Path search: {sys.path}') + shared.log.error("Module not found: k-diffusion") + sys.exit(1) + + samplers_k_diffusion = [ ('Euler', 'sample_euler', ['k_euler'], {"scheduler": "default"}), ('Euler a', 'sample_euler_ancestral', ['k_euler_a', 'k_euler_ancestral'], {"scheduler": "default", "brownian_noise": False}), @@ -32,7 +51,7 @@ samplers_k_diffusion = [ samplers_data_k_diffusion = [ sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options) for label, funcname, aliases, options in samplers_k_diffusion - if hasattr(k_diffusion.sampling, funcname) + if hasattr(k_sampling, funcname) ] sampler_extra_params = { @@ -79,7 +98,6 @@ class CFGDenoiser(torch.nn.Module): while shared.state.paused: if shared.state.interrupted or shared.state.skipped: raise sd_samplers_common.InterruptedException - import time time.sleep(0.1) # at self.image_cfg_scale == 1.0 produced results for edit model are the same as with normal sampling, # so is_edit_model is set to False to support AND composition. @@ -214,7 +232,7 @@ class KDiffusionSampler: denoiser = k_diffusion.external.CompVisVDenoiser if sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization) self.funcname = funcname - self.func = getattr(k_diffusion.sampling, self.funcname) + self.func = getattr(k_sampling, self.funcname) self.extra_params = sampler_extra_params.get(funcname, []) self.model_wrap_cfg = CFGDenoiser(self.model_wrap) self.sampler_noises = None @@ -258,7 +276,7 @@ class KDiffusionSampler: self.model_wrap_cfg.image_cfg_scale = getattr(p, 'image_cfg_scale', None) self.eta = p.eta if p.eta is not None else shared.opts.scheduler_eta self.s_min_uncond = getattr(p, 's_min_uncond', 0.0) - k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else []) + k_sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else []) extra_params_kwargs = {} for param_name in self.extra_params: if hasattr(p, param_name) and param_name in inspect.signature(self.func).parameters: @@ -277,17 +295,17 @@ class KDiffusionSampler: elif self.config.options.get('scheduler', None) == 'karras': sigma_min = p.s_min if p.s_min > 0 else self.model_wrap.sigmas[0].item() sigma_max = p.s_max if p.s_max > 0 else self.model_wrap.sigmas[-1].item() - sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) + sigmas = k_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) elif self.config.options.get('scheduler', None) == 'exponential': sigma_min = p.s_min if p.s_min > 0 else self.model_wrap.sigmas[0].item() sigma_max = p.s_max if p.s_max > 0 else self.model_wrap.sigmas[-1].item() - sigmas = k_diffusion.sampling.get_sigmas_exponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) + sigmas = k_sampling.get_sigmas_exponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) elif self.config.options.get('scheduler', None) == 'polyexponential': sigma_min = p.s_min if p.s_min > 0 else self.model_wrap.sigmas[0].item() sigma_max = p.s_max if p.s_max > 0 else self.model_wrap.sigmas[-1].item() - sigmas = k_diffusion.sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) + sigmas = k_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, device=shared.device) elif self.config.options.get('scheduler', None) == 'vp': - sigmas = k_diffusion.sampling.get_sigmas_vp(n=steps, device=shared.device) + sigmas = k_sampling.get_sigmas_vp(n=steps, device=shared.device) if discard_next_to_last_sigma: sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) return sigmas @@ -296,7 +314,6 @@ class KDiffusionSampler: """For DPM++ SDE: manually create noise sampler to enable deterministic results across different batch sizes""" if shared.opts.no_dpmpp_sde_batch_determinism: return None - from k_diffusion.sampling import BrownianTreeNoiseSampler positive_sigmas = sigmas[sigmas > 0] if positive_sigmas.numel() > 0: sigma_min = positive_sigmas.min(dim=0)[0] @@ -304,7 +321,7 @@ class KDiffusionSampler: sigma_min = 0 sigma_max = sigmas.max() current_iter_seeds = p.all_seeds[p.iteration * p.batch_size:(p.iteration + 1) * p.batch_size] - return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds) + return k_sampling.BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds) def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None): steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps) diff --git a/modules/shared.py b/modules/shared.py index 1ff76162c..33a309659 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -23,7 +23,7 @@ from installer import print_dict from installer import log as central_logger # pylint: disable=E0611 -errors.install(gr) +errors.install([gr]) demo: gr.Blocks = None log = central_logger progress_print_out = sys.stdout