a1111 compatibility items

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:

modules/sd_samplers_cfg_denoiser.py

@@ -1,3 +1,4 @@
+# TODO a1111 compatibility module
 # TODO cfg_denoiser implementation missing
 
 import torch

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

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

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:

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