automatic/modules/models/diffusion/uni_pc/sampler.py

"""SAMPLING ONLY."""

import numpy as np
import torch

from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
from modules import shared, devices
from ldm.modules.diffusionmodules.util import extract_into_tensor


class UniPCSampler(object):
    def __init__(self, model, **kwargs):
        super().__init__()
        self.model = model
        to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
        self.before_sample = None
        self.after_sample = None
        self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))

    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
        # persist steps so we can eventually find denoising strength
        self.inflated_steps = ddim_num_steps

    @torch.no_grad()
    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
        if noise is None:
            noise = torch.randn_like(x0)

        # first time we have all the info to get the real parameters from the ui
        # value from the hires steps slider:
        num_inference_steps = t[0] + 1
        # (num_inference_steps // denoising_strength):
        inflated_steps = self.inflated_steps
        # not exact:
        self.denoising_strength = num_inference_steps/inflated_steps

        # values used for timesteps that generate noise in diffusers repo
        init_timestep = min(
                int(num_inference_steps * self.denoising_strength),
                num_inference_steps,
            )
        t_start = max(num_inference_steps - init_timestep, 0)

        # actual number of steps we'll run
        self.steps = max(
                num_inference_steps - init_timestep,
                shared.opts.uni_pc_order+1,
            )

        t = torch.full(t.shape, self.steps).to(t.device)

        scheduler_timesteps = np.linspace(
                0,
                self.model.num_timesteps-1,
                num_inference_steps + 1,
            ).round()[::-1][:-1].copy().astype(np.int64)
        _, unique_indices = np.unique(scheduler_timesteps, return_index=True)
        scheduler_timesteps = scheduler_timesteps[np.sort(unique_indices)]
        scheduler_timesteps = torch.from_numpy(scheduler_timesteps).to(t.device)

        sample_timesteps = scheduler_timesteps[t_start:]
        latent_timestep = sample_timesteps[:1].repeat(x0.shape[0])

        alphas_cumprod = self.alphas_cumprod
        sqrt_alphas_prod = alphas_cumprod[latent_timestep] ** 0.5
        sqrt_alphas_prod = sqrt_alphas_prod.flatten()

        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[latent_timestep]) ** 0.5
        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()

        return (sqrt_alphas_prod * x0 + sqrt_one_minus_alpha_prod * noise)

    def decode(self, x_latent, conditioning, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
               use_original_steps=False, callback=None):
        #print(f'steps {self.steps} denoising {self.denoising_strength}')

        noise_schedule = NoiseScheduleVP("discrete", alphas_cumprod=self.alphas_cumprod)

        # same as in .sample(), i guess
        model_type = "v" if self.model.parameterization == "v" else "noise"

        model_fn = model_wrapper(
            lambda x, t, c: self.model.apply_model(x, t, c),
            noise_schedule,
            model_type=model_type,
            guidance_type="classifier-free",
            #condition=conditioning,
            #unconditional_condition=unconditional_conditioning,
            guidance_scale=unconditional_guidance_scale,
        )

        self.uni_pc = UniPC(
                model_fn,
                noise_schedule,
                predict_x0=True,
                thresholding=False,
                variant=shared.opts.uni_pc_variant,
                condition=conditioning,
                unconditional_condition=unconditional_conditioning,
                before_sample=self.before_sample,
                after_sample=self.after_sample,
                after_update=self.after_update,
            )

        return self.uni_pc.sample(
                x_latent,
                steps=self.steps,
                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,
                t_start=self.denoising_strength,
            )

    def register_buffer(self, name, attr):
        if type(attr) == torch.Tensor:
            if attr.device != devices.device:
                attr = attr.to(devices.device)
        setattr(self, name, attr)

    def set_hooks(self, before_sample, after_sample, after_update):
        self.before_sample = before_sample
        self.after_sample = after_sample
        self.after_update = after_update

    @torch.no_grad()
    def sample(self,
               S,
               batch_size,
               shape,
               conditioning=None,
               callback=None,
               normals_sequence=None,
               img_callback=None,
               quantize_x0=False,
               eta=0.,
               mask=None,
               x0=None,
               temperature=1.,
               noise_dropout=0.,
               score_corrector=None,
               corrector_kwargs=None,
               verbose=True,
               x_T=None,
               log_every_t=100,
               unconditional_guidance_scale=1.,
               unconditional_conditioning=None,
               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
               **kwargs
               ):
        if conditioning is not None:
            if isinstance(conditioning, dict):
                ctmp = conditioning[list(conditioning.keys())[0]]
                while isinstance(ctmp, list): ctmp = ctmp[0]
                cbs = ctmp.shape[0]
                if cbs != batch_size:
                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")

            elif isinstance(conditioning, list):
                for ctmp in conditioning:
                    if ctmp.shape[0] != batch_size:
                        print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")

            else:
                if conditioning.shape[0] != batch_size:
                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")

        # sampling
        C, H, W = shape
        size = (batch_size, C, H, W)
        # print(f'Data shape for UniPC sampling is {size}')

        device = self.model.betas.device
        if x_T is None:
            img = torch.randn(size, device=device)
        else:
            img = x_T

        ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)

        # SD 1.X is "noise", SD 2.X is "v"
        model_type = "v" if self.model.parameterization == "v" else "noise"

        model_fn = model_wrapper(
            lambda x, t, c: self.model.apply_model(x, t, c),
            ns,
            model_type=model_type,
            guidance_type="classifier-free",
            #condition=conditioning,
            #unconditional_condition=unconditional_conditioning,
            guidance_scale=unconditional_guidance_scale,
        )

        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update)
        x = uni_pc.sample(img, steps=S, 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.to(device), None