From c77bd7dcef7db53759da7c48e0ca9ad389d289c2 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 17 Oct 2023 09:22:40 -0400 Subject: [PATCH] a1111 compatibility items --- installer.py | 1 + modules/sd_samplers_cfg_denoiser.py | 1 + modules/sd_samplers_timesteps.py | 72 +++++++++++++ modules/sd_samplers_timesteps_impl.py | 139 ++++++++++++++++++++++++++ modules/shared.py | 120 +--------------------- modules/shared_state.py | 122 ++++++++++++++++++++++ 6 files changed, 337 insertions(+), 118 deletions(-) create mode 100644 modules/sd_samplers_timesteps.py create mode 100644 modules/sd_samplers_timesteps_impl.py create mode 100644 modules/shared_state.py diff --git a/installer.py b/installer.py index c1493b093..36121cdba 100644 --- a/installer.py +++ b/installer.py @@ -447,6 +447,7 @@ def check_torch(): log.info('Using CPU-only Torch') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') if 'torch' in torch_command and not args.version: + log.info('Installing torch - this may take a while') install(torch_command, 'torch torchvision') else: try: diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py index c80a80e50..8de26b8f3 100644 --- a/modules/sd_samplers_cfg_denoiser.py +++ b/modules/sd_samplers_cfg_denoiser.py @@ -1,3 +1,4 @@ +# TODO a1111 compatibility module # TODO cfg_denoiser implementation missing import torch diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py new file mode 100644 index 000000000..dfe0a857e --- /dev/null +++ b/modules/sd_samplers_timesteps.py @@ -0,0 +1,72 @@ +# TODO a1111 compatibility module + +import torch +from modules import sd_samplers_common, sd_samplers_timesteps_impl, sd_samplers_compvis +from modules.sd_samplers_cfg_denoiser import CFGDenoiser +import modules.shared as shared + + +samplers_timesteps = [ + ('DDIM', sd_samplers_timesteps_impl.ddim, ['ddim'], {}), + ('PLMS', sd_samplers_timesteps_impl.plms, ['plms'], {}), + ('UniPC', sd_samplers_timesteps_impl.unipc, ['unipc'], {}), +] + + +samplers_data_timesteps = [ + sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: VanillaStableDiffusionSampler(funcname, model), aliases, options) + for label, funcname, aliases, options in samplers_timesteps +] + + +class CompVisTimestepsDenoiser(torch.nn.Module): + def __init__(self, model, *args, **kwargs): + super().__init__(*args, **kwargs) + self.inner_model = model + + def forward(self, input, timesteps, **kwargs): # pylint: disable=redefined-builtin + return self.inner_model.apply_model(input, timesteps, **kwargs) + + +class CompVisTimestepsVDenoiser(torch.nn.Module): + def __init__(self, model, *args, **kwargs): + super().__init__(*args, **kwargs) + self.inner_model = model + + def predict_eps_from_z_and_v(self, x_t, t, v): + return self.inner_model.sqrt_alphas_cumprod[t.to(torch.int), None, None, None] * v + self.inner_model.sqrt_one_minus_alphas_cumprod[t.to(torch.int), None, None, None] * x_t + + def forward(self, input, timesteps, **kwargs): # pylint: disable=redefined-builtin + model_output = self.inner_model.apply_model(input, timesteps, **kwargs) + e_t = self.predict_eps_from_z_and_v(input, timesteps, model_output) + return e_t + + +class CFGDenoiserTimesteps(CFGDenoiser): + + def __init__(self, sampler): + super().__init__(sampler) + + self.alphas = shared.sd_model.alphas_cumprod + self.mask_before_denoising = True + + def get_pred_x0(self, x_in, x_out, sigma): + ts = sigma.to(dtype=int) + + a_t = self.alphas[ts][:, None, None, None] + sqrt_one_minus_at = (1 - a_t).sqrt() + + pred_x0 = (x_in - sqrt_one_minus_at * x_out) / a_t.sqrt() + + return pred_x0 + + @property + def inner_model(self): + if self.model_wrap is None: + denoiser = CompVisTimestepsVDenoiser if shared.sd_model.parameterization == "v" else CompVisTimestepsDenoiser + self.model_wrap = denoiser(shared.sd_model) + + return self.model_wrap + + +VanillaStableDiffusionSampler = sd_samplers_compvis.VanillaStableDiffusionSampler diff --git a/modules/sd_samplers_timesteps_impl.py b/modules/sd_samplers_timesteps_impl.py new file mode 100644 index 000000000..a15c76a4a --- /dev/null +++ b/modules/sd_samplers_timesteps_impl.py @@ -0,0 +1,139 @@ +# TODO a1111 compatibility module + +import torch +import tqdm +import k_diffusion.sampling +import numpy as np + +from modules import shared +from modules.unipc import uni_pc + + +@torch.no_grad() +def ddim(model, x, timesteps, extra_args=None, callback=None, disable=None, eta=0.0): + alphas_cumprod = model.inner_model.inner_model.alphas_cumprod + alphas = alphas_cumprod[timesteps] + alphas_prev = alphas_cumprod[torch.nn.functional.pad(timesteps[:-1], pad=(1, 0))].to(torch.float64 if x.device.type != 'mps' else torch.float32) # pylint: disable=not-callable + sqrt_one_minus_alphas = torch.sqrt(1 - alphas) + sigmas = eta * np.sqrt((1 - alphas_prev.cpu().numpy()) / (1 - alphas.cpu()) * (1 - alphas.cpu() / alphas_prev.cpu().numpy())) + + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones((x.shape[0])) + s_x = x.new_ones((x.shape[0], 1, 1, 1)) + for i in tqdm.trange(len(timesteps) - 1, disable=disable): + index = len(timesteps) - 1 - i + + e_t = model(x, timesteps[index].item() * s_in, **extra_args) + + a_t = alphas[index].item() * s_x + a_prev = alphas_prev[index].item() * s_x + sigma_t = sigmas[index].item() * s_x + sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x + + pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() + dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t + noise = sigma_t * k_diffusion.sampling.torch.randn_like(x) + x = a_prev.sqrt() * pred_x0 + dir_xt + noise + + if callback is not None: + callback({'x': x, 'i': i, 'sigma': 0, 'sigma_hat': 0, 'denoised': pred_x0}) + + return x + + +@torch.no_grad() +def plms(model, x, timesteps, extra_args=None, callback=None, disable=None): + alphas_cumprod = model.inner_model.inner_model.alphas_cumprod + alphas = alphas_cumprod[timesteps] + alphas_prev = alphas_cumprod[torch.nn.functional.pad(timesteps[:-1], pad=(1, 0))].to(torch.float64 if x.device.type != 'mps' else torch.float32) # pylint: disable=not-callable + sqrt_one_minus_alphas = torch.sqrt(1 - alphas) + + extra_args = {} if extra_args is None else extra_args + s_in = x.new_ones([x.shape[0]]) + s_x = x.new_ones((x.shape[0], 1, 1, 1)) + old_eps = [] + + def get_x_prev_and_pred_x0(e_t, index): + # select parameters corresponding to the currently considered timestep + a_t = alphas[index].item() * s_x + a_prev = alphas_prev[index].item() * s_x + sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x + + # current prediction for x_0 + pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() + + # direction pointing to x_t + dir_xt = (1. - a_prev).sqrt() * e_t + x_prev = a_prev.sqrt() * pred_x0 + dir_xt + return x_prev, pred_x0 + + for i in tqdm.trange(len(timesteps) - 1, disable=disable): + index = len(timesteps) - 1 - i + ts = timesteps[index].item() * s_in + t_next = timesteps[max(index - 1, 0)].item() * s_in + + e_t = model(x, ts, **extra_args) + + if len(old_eps) == 0: + # Pseudo Improved Euler (2nd order) + x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index) + e_t_next = model(x_prev, t_next, **extra_args) + e_t_prime = (e_t + e_t_next) / 2 + elif len(old_eps) == 1: + # 2nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (3 * e_t - old_eps[-1]) / 2 + elif len(old_eps) == 2: + # 3nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12 + else: + # 4nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24 + + x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index) + + old_eps.append(e_t) + if len(old_eps) >= 4: + old_eps.pop(0) + + x = x_prev + + if callback is not None: + callback({'x': x, 'i': i, 'sigma': 0, 'sigma_hat': 0, 'denoised': pred_x0}) + + return x + + +class UniPCCFG(uni_pc.UniPC): + def __init__(self, cfg_model, extra_args, callback, *args, **kwargs): + super().__init__(None, *args, **kwargs) + + def after_update(x, model_x): + callback({'x': x, 'i': self.index, 'sigma': 0, 'sigma_hat': 0, 'denoised': model_x}) + self.index += 1 + + self.cfg_model = cfg_model + self.extra_args = extra_args + self.callback = callback + self.index = 0 + self.after_update = after_update + + def get_model_input_time(self, t_continuous): + return (t_continuous - 1. / self.noise_schedule.total_N) * 1000. + + def model(self, x, t): + t_input = self.get_model_input_time(t) + + res = self.cfg_model(x, t_input, **self.extra_args) + + return res + + +def unipc(model, x, timesteps, extra_args=None, callback=None, disable=None, is_img2img=False): # pylint: disable=unused-argument + alphas_cumprod = model.inner_model.inner_model.alphas_cumprod + + ns = uni_pc.NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod) + t_start = timesteps[-1] / 1000 + 1 / 1000 if is_img2img else None # this is likely off by a bit - if someone wants to fix it please by all means + unipc_sampler = UniPCCFG(model, extra_args, callback, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant) + x = unipc_sampler.sample(x, steps=len(timesteps), t_start=t_start, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final) + + return x diff --git a/modules/shared.py b/modules/shared.py index 9dc7e7720..6ceed6b12 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -3,7 +3,6 @@ import os import sys import time import json -import datetime import contextlib import urllib.request from types import SimpleNamespace @@ -13,7 +12,7 @@ import requests import gradio as gr import fasteners from rich.console import Console -from modules import errors, shared_items, cmd_args, ui_components +from modules import errors, shared_items, shared_state, cmd_args, ui_components from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack import modules.interrogate @@ -75,122 +74,7 @@ class Backend(Enum): DIFFUSERS = 2 - -class State: - skipped = False - interrupted = False - paused = False - job = "" - job_no = 0 - job_count = 0 - total_jobs = 0 - processing_has_refined_job_count = False - job_timestamp = '0' - sampling_step = 0 - sampling_steps = 0 - current_latent = None - current_image = None - current_image_sampling_step = 0 - id_live_preview = 0 - textinfo = None - time_start = None - need_restart = False - server_start = time.time() - oom = False - debug_output = os.environ.get('SD_STATE_DEBUG', None) - - def skip(self): - log.debug('Requested skip') - self.skipped = True - - def interrupt(self): - log.debug('Requested interrupt') - self.interrupted = True - - def pause(self): - self.paused = not self.paused - log.debug(f'Requested {"pause" if self.paused else "continue"}') - - def nextjob(self): - if opts.live_previews_enable and opts.show_progress_every_n_steps == -1: - self.do_set_current_image() - self.job_no += 1 - self.sampling_step = 0 - self.current_image_sampling_step = 0 - - def dict(self): - obj = { - "skipped": self.skipped, - "interrupted": self.interrupted, - "job": self.job, - "job_count": self.job_count, - "job_timestamp": self.job_timestamp, - "job_no": self.job_no, - "sampling_step": self.sampling_step, - "sampling_steps": self.sampling_steps, - } - return obj - - def begin(self, title=""): - self.total_jobs += 1 - self.current_image = None - self.current_image_sampling_step = 0 - self.current_latent = None - self.id_live_preview = 0 - self.interrupted = False - self.job = title - self.job_count = -1 - self.job_no = 0 - self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") - self.paused = False - self.processing_has_refined_job_count = False - self.sampling_step = 0 - self.skipped = False - self.textinfo = None - self.time_start = time.time() - if self.debug_output: - log.debug(f'State begin: {self.job}') - devices.torch_gc() - - def end(self): - if self.time_start is None: # someone called end before being - log.debug(f'Access state.end: {sys._getframe().f_back.f_code.co_name}') # pylint: disable=protected-access - self.time_start = time.time() - if self.debug_output: - log.debug(f'State end: {self.job} time={time.time() - self.time_start:.2f}s') - self.job = "" - self.job_count = 0 - self.job_no = 0 - self.paused = False - self.interrupted = False - self.skipped = False - devices.torch_gc() - - def set_current_image(self): - """sets self.current_image from self.current_latent if enough sampling steps have been made after the last call to this""" - if not parallel_processing_allowed: - return - if abs(self.sampling_step - self.current_image_sampling_step) >= opts.show_progress_every_n_steps and opts.live_previews_enable and opts.show_progress_every_n_steps > 0: - self.do_set_current_image() - - def do_set_current_image(self): - if self.current_latent is None: - return - import modules.sd_samplers # pylint: disable=W0621 - try: - image = modules.sd_samplers.samples_to_image_grid(self.current_latent) if opts.show_progress_grid else modules.sd_samplers.sample_to_image(self.current_latent) - self.assign_current_image(image) - self.current_image_sampling_step = self.sampling_step - except Exception: - # log.error(f'Error setting current image: step={self.sampling_step} {e}') - pass - - def assign_current_image(self, image): - self.current_image = image - self.id_live_preview += 1 - - -state = State() +state = shared_state.State() if not hasattr(cmd_opts, "use_openvino"): cmd_opts.use_openvino = False if cmd_opts.use_openvino: diff --git a/modules/shared_state.py b/modules/shared_state.py new file mode 100644 index 000000000..ad2d992d1 --- /dev/null +++ b/modules/shared_state.py @@ -0,0 +1,122 @@ +import os +import sys +import time +import datetime +from modules.errors import log + + +class State: + skipped = False + interrupted = False + paused = False + job = "" + job_no = 0 + job_count = 0 + total_jobs = 0 + processing_has_refined_job_count = False + job_timestamp = '0' + sampling_step = 0 + sampling_steps = 0 + current_latent = None + current_image = None + current_image_sampling_step = 0 + id_live_preview = 0 + textinfo = None + time_start = None + need_restart = False + server_start = time.time() + oom = False + debug_output = os.environ.get('SD_STATE_DEBUG', None) + + def skip(self): + log.debug('Requested skip') + self.skipped = True + + def interrupt(self): + log.debug('Requested interrupt') + self.interrupted = True + + def pause(self): + self.paused = not self.paused + log.debug(f'Requested {"pause" if self.paused else "continue"}') + + def nextjob(self): + self.do_set_current_image() + self.job_no += 1 + self.sampling_step = 0 + self.current_image_sampling_step = 0 + + def dict(self): + obj = { + "skipped": self.skipped, + "interrupted": self.interrupted, + "job": self.job, + "job_count": self.job_count, + "job_timestamp": self.job_timestamp, + "job_no": self.job_no, + "sampling_step": self.sampling_step, + "sampling_steps": self.sampling_steps, + } + return obj + + def begin(self, title=""): + import modules.devices + self.total_jobs += 1 + self.current_image = None + self.current_image_sampling_step = 0 + self.current_latent = None + self.id_live_preview = 0 + self.interrupted = False + self.job = title + self.job_count = -1 + self.job_no = 0 + self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + self.paused = False + self.processing_has_refined_job_count = False + self.sampling_step = 0 + self.skipped = False + self.textinfo = None + self.time_start = time.time() + if self.debug_output: + log.debug(f'State begin: {self.job}') + modules.devices.torch_gc() + + def end(self): + import modules.devices + if self.time_start is None: # someone called end before being + log.debug(f'Access state.end: {sys._getframe().f_back.f_code.co_name}') # pylint: disable=protected-access + self.time_start = time.time() + if self.debug_output: + log.debug(f'State end: {self.job} time={time.time() - self.time_start:.2f}s') + self.job = "" + self.job_count = 0 + self.job_no = 0 + self.paused = False + self.interrupted = False + self.skipped = False + modules.devices.torch_gc() + + def set_current_image(self): + from modules.shared import opts, cmd_opts + """sets self.current_image from self.current_latent if enough sampling steps have been made after the last call to this""" + if cmd_opts.lowvram: + return + if abs(self.sampling_step - self.current_image_sampling_step) >= opts.show_progress_every_n_steps and opts.live_previews_enable and opts.show_progress_every_n_steps > 0: + self.do_set_current_image() + + def do_set_current_image(self): + from modules.shared import opts + if self.current_latent is None: + return + import modules.sd_samplers # pylint: disable=W0621 + try: + image = modules.sd_samplers.samples_to_image_grid(self.current_latent) if opts.show_progress_grid else modules.sd_samplers.sample_to_image(self.current_latent) + self.assign_current_image(image) + self.current_image_sampling_step = self.sampling_step + except Exception: + # log.error(f'Error setting current image: step={self.sampling_step} {e}') + pass + + def assign_current_image(self, image): + self.current_image = image + self.id_live_preview += 1