# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
Based on: https://github.com/crowsonkb/k-diffusion
"""
import random

import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F

# from piq import LPIPS
from utils.ssim import SSIM

from modules.diffusion.karras.random_utils import get_generator


def mean_flat(tensor):
    """
    Take the mean over all non-batch dimensions.
    """
    return tensor.mean(dim=list(range(1, len(tensor.shape))))


def append_dims(x, target_dims):
    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
    dims_to_append = target_dims - x.ndim
    if dims_to_append < 0:
        raise ValueError(
            f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
        )
    return x[(...,) + (None,) * dims_to_append]


def append_zero(x):
    return th.cat([x, x.new_zeros([1])])


def get_weightings(weight_schedule, snrs, sigma_data):
    if weight_schedule == "snr":
        weightings = snrs
    elif weight_schedule == "snr+1":
        weightings = snrs + 1
    elif weight_schedule == "karras":
        weightings = snrs + 1.0 / sigma_data**2
    elif weight_schedule == "truncated-snr":
        weightings = th.clamp(snrs, min=1.0)
    elif weight_schedule == "uniform":
        weightings = th.ones_like(snrs)
    else:
        raise NotImplementedError()
    return weightings


class KarrasDenoiser:
    def __init__(
        self,
        sigma_data: float = 0.5,
        sigma_max=80.0,
        sigma_min=0.002,
        rho=7.0,
        weight_schedule="karras",
        distillation=False,
        loss_norm="l2",
    ):
        self.sigma_data = sigma_data
        self.sigma_max = sigma_max
        self.sigma_min = sigma_min
        self.weight_schedule = weight_schedule
        self.distillation = distillation
        self.loss_norm = loss_norm
        # if loss_norm == "lpips":
        #     self.lpips_loss = LPIPS(replace_pooling=True, reduction="none")
        if loss_norm == "ssim":
            self.ssim_loss = SSIM()
        self.rho = rho
        self.num_timesteps = 40

    def get_snr(self, sigmas):
        return sigmas**-2

    def get_sigmas(self, sigmas):
        return sigmas

    def get_scalings(self, sigma):
        c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
        c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5
        c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
        return c_skip, c_out, c_in

    def get_scalings_for_boundary_condition(self, sigma):
        c_skip = self.sigma_data**2 / (
            (sigma - self.sigma_min) ** 2 + self.sigma_data**2
        )
        c_out = (
            (sigma - self.sigma_min)
            * self.sigma_data
            / (sigma**2 + self.sigma_data**2) ** 0.5
        )
        c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
        return c_skip, c_out, c_in

    def training_losses(self, model, x_start, sigmas, condition=None, noise=None):
        if noise is None:
            noise = th.randn_like(x_start)

        terms = {}

        dims = x_start.ndim
        x_t = x_start + noise * append_dims(sigmas, dims)
        model_output, denoised = self.denoise(model, x_t, sigmas, condition)

        snrs = self.get_snr(sigmas)
        weights = append_dims(
            get_weightings(self.weight_schedule, snrs, self.sigma_data), dims
        )
        # terms["xs_mse"] = mean_flat((denoised - x_start) ** 2)
        terms["mse"] = mean_flat(weights * (denoised - x_start) ** 2)
        # terms["mae"] = mean_flat(weights * th.abs(denoised - x_start))
        # terms["mse"] = nn.MSELoss(reduction="none")(denoised, x_start)

        # if "vb" in terms:
        #     terms["loss"] = terms["mse"] + terms["vb"]
        # else:
        terms["loss"] = terms["mse"]

        return terms

    def consistency_losses(
        self,
        model,
        x_start,
        num_scales,
        # model_kwargs=None,
        condition=None,
        target_model=None,
        teacher_model=None,
        teacher_diffusion=None,
        noise=None,
    ):
        if noise is None:
            noise = th.randn_like(x_start)

        dims = x_start.ndim

        def denoise_fn(x, t):
            return self.denoise(model, x, t, condition)[1]

        if target_model:

            @th.no_grad()
            def target_denoise_fn(x, t):
                return self.denoise(target_model, x, t, condition)[1]

        else:
            raise NotImplementedError("Must have a target model")

        if teacher_model:

            @th.no_grad()
            def teacher_denoise_fn(x, t):
                return teacher_diffusion.denoise(teacher_model, x, t, condition)[1]

        @th.no_grad()
        def heun_solver(samples, t, next_t, x0):
            x = samples
            if teacher_model is None:
                denoiser = x0
            else:
                denoiser = teacher_denoise_fn(x, t)

            d = (x - denoiser) / append_dims(t, dims)
            samples = x + d * append_dims(next_t - t, dims)
            if teacher_model is None:
                denoiser = x0
            else:
                denoiser = teacher_denoise_fn(samples, next_t)

            next_d = (samples - denoiser) / append_dims(next_t, dims)
            samples = x + (d + next_d) * append_dims((next_t - t) / 2, dims)

            return samples

        @th.no_grad()
        def euler_solver(samples, t, next_t, x0):
            x = samples
            if teacher_model is None:
                denoiser = x0
            else:
                denoiser = teacher_denoise_fn(x, t)
            d = (x - denoiser) / append_dims(t, dims)
            samples = x + d * append_dims(next_t - t, dims)

            return samples

        indices = th.randint(
            0, num_scales - 1, (x_start.shape[0],), device=x_start.device
        )

        t = self.sigma_max ** (1 / self.rho) + indices / (num_scales - 1) * (
            self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
        )
        t = t**self.rho

        t2 = self.sigma_max ** (1 / self.rho) + (indices + 1) / (num_scales - 1) * (
            self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
        )
        t2 = t2**self.rho

        x_t = x_start + noise * append_dims(t, dims)

        dropout_state = th.get_rng_state()
        distiller = denoise_fn(x_t, t)

        if teacher_model is None:
            x_t2 = euler_solver(x_t, t, t2, x_start).detach()
        else:
            x_t2 = heun_solver(x_t, t, t2, x_start).detach()

        th.set_rng_state(dropout_state)
        distiller_target = target_denoise_fn(x_t2, t2)
        distiller_target = distiller_target.detach()

        snrs = self.get_snr(t)
        weights = get_weightings(self.weight_schedule, snrs, self.sigma_data)
        if self.loss_norm == "l1":
            diffs = th.abs(distiller - distiller_target)
            loss = mean_flat(diffs) * weights
        elif self.loss_norm == "l2":
            # diffs = (distiller - distiller_target) ** 2
            loss = F.mse_loss(distiller, distiller_target)
            # loss = mean_flat(diffs) * weights
        elif self.loss_norm == "ssim":
            loss = self.ssim_loss(distiller, distiller_target) * weights
        # elif self.loss_norm == "l2-32":
        #     distiller = F.interpolate(distiller, size=32, mode="bilinear")
        #     distiller_target = F.interpolate(
        #         distiller_target,
        #         size=32,
        #         mode="bilinear",
        #     )
        #     diffs = (distiller - distiller_target) ** 2
        #     loss = mean_flat(diffs) * weights
        # elif self.loss_norm == "lpips":
        #     if x_start.shape[-1] < 256:
        #         distiller = F.interpolate(distiller, size=224, mode="bilinear")
        #         distiller_target = F.interpolate(
        #             distiller_target, size=224, mode="bilinear"
        #         )

        #     loss = (
        #         self.lpips_loss(
        #             (distiller + 1) / 2.0,
        #             (distiller_target + 1) / 2.0,
        #         )
        #         * weights
        #     )
        else:
            raise ValueError(f"Unknown loss norm {self.loss_norm}")

        terms = {}
        terms["loss"] = loss

        return terms

    # def progdist_losses(
    #     self,
    #     model,
    #     x_start,
    #     num_scales,
    #     model_kwargs=None,
    #     teacher_model=None,
    #     teacher_diffusion=None,
    #     noise=None,
    # ):
    #     if model_kwargs is None:
    #         model_kwargs = {}
    #     if noise is None:
    #         noise = th.randn_like(x_start)

    #     dims = x_start.ndim

    #     def denoise_fn(x, t):
    #         return self.denoise(model, x, t, **model_kwargs)[1]

    #     @th.no_grad()
    #     def teacher_denoise_fn(x, t):
    #         return teacher_diffusion.denoise(teacher_model, x, t, **model_kwargs)[1]

    #     @th.no_grad()
    #     def euler_solver(samples, t, next_t):
    #         x = samples
    #         denoiser = teacher_denoise_fn(x, t)
    #         d = (x - denoiser) / append_dims(t, dims)
    #         samples = x + d * append_dims(next_t - t, dims)

    #         return samples

    #     @th.no_grad()
    #     def euler_to_denoiser(x_t, t, x_next_t, next_t):
    #         denoiser = x_t - append_dims(t, dims) * (x_next_t - x_t) / append_dims(
    #             next_t - t, dims
    #         )
    #         return denoiser

    #     indices = th.randint(0, num_scales, (x_start.shape[0],), device=x_start.device)

    #     t = self.sigma_max ** (1 / self.rho) + indices / num_scales * (
    #         self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
    #     )
    #     t = t**self.rho

    #     t2 = self.sigma_max ** (1 / self.rho) + (indices + 0.5) / num_scales * (
    #         self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
    #     )
    #     t2 = t2**self.rho

    #     t3 = self.sigma_max ** (1 / self.rho) + (indices + 1) / num_scales * (
    #         self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho)
    #     )
    #     t3 = t3**self.rho

    #     x_t = x_start + noise * append_dims(t, dims)

    #     denoised_x = denoise_fn(x_t, t)

    #     x_t2 = euler_solver(x_t, t, t2).detach()
    #     x_t3 = euler_solver(x_t2, t2, t3).detach()

    #     target_x = euler_to_denoiser(x_t, t, x_t3, t3).detach()

    #     snrs = self.get_snr(t)
    #     weights = get_weightings(self.weight_schedule, snrs, self.sigma_data)
    #     if self.loss_norm == "l1":
    #         diffs = th.abs(denoised_x - target_x)
    #         loss = mean_flat(diffs) * weights
    #     elif self.loss_norm == "l2":
    #         diffs = (denoised_x - target_x) ** 2
    #         loss = mean_flat(diffs) * weights
    #     elif self.loss_norm == "lpips":
    #         if x_start.shape[-1] < 256:
    #             denoised_x = F.interpolate(denoised_x, size=224, mode="bilinear")
    #             target_x = F.interpolate(target_x, size=224, mode="bilinear")
    #         loss = (
    #             self.lpips_loss(
    #                 (denoised_x + 1) / 2.0,
    #                 (target_x + 1) / 2.0,
    #             )
    #             * weights
    #         )
    #     else:
    #         raise ValueError(f"Unknown loss norm {self.loss_norm}")

    #     terms = {}
    #     terms["loss"] = loss

    #     return terms

    def denoise(self, model, x_t, sigmas, condition):
        if not self.distillation:
            c_skip, c_out, c_in = [
                append_dims(x, x_t.ndim) for x in self.get_scalings(sigmas)
            ]
        else:
            c_skip, c_out, c_in = [
                append_dims(x, x_t.ndim)
                for x in self.get_scalings_for_boundary_condition(sigmas)
            ]
        rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
        # rescaled_t = rescaled_t[:, None]
        model_output = model(c_in * x_t, rescaled_t, condition)
        denoised = c_out * model_output + c_skip * x_t
        return model_output, denoised


def karras_sample(
    diffusion,
    model,
    shape,
    steps,
    clip_denoised=True,
    progress=True,
    callback=None,
    # model_kwargs=None,
    condition=None,
    device=None,
    sigma_min=0.002,
    sigma_max=80,  # higher for highres?
    rho=7.0,
    sampler="heun",
    s_churn=0.0,
    s_tmin=0.0,
    s_tmax=float("inf"),
    s_noise=1.0,
    generator=None,
    ts=None,
):
    if generator is None:
        generator = get_generator("dummy")

    if sampler == "progdist":
        sigmas = get_sigmas_karras(steps + 1, sigma_min, sigma_max, rho, device=device)
    else:
        sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
    th.manual_seed(42)
    x_T = generator.randn(*shape, device=device) * sigma_max
    sigmas = sigmas.unsqueeze(-1)
    sample_fn = {
        "heun": sample_heun,
        "dpm": sample_dpm,
        "ancestral": sample_euler_ancestral,
        "onestep": sample_onestep,
        "progdist": sample_progdist,
        "euler": sample_euler,
        "multistep": stochastic_iterative_sampler,
    }[sampler]

    if sampler in ["heun", "dpm"]:
        sampler_args = dict(
            s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise
        )
    elif sampler == "multistep":
        sampler_args = dict(
            ts=ts, t_min=sigma_min, t_max=sigma_max, rho=diffusion.rho, steps=steps
        )
    else:
        sampler_args = {}

    def denoiser(x_t, sigma):
        _, denoised = diffusion.denoise(model, x_t, sigma, condition)
        if clip_denoised:
            denoised = denoised.clamp(-1, 1)
        return denoised

    x_0 = sample_fn(
        denoiser,
        x_T,
        sigmas,
        generator,
        progress=progress,
        callback=callback,
        **sampler_args,
    )
    return x_0.clamp(-1, 1)


def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu"):
    """Constructs the noise schedule of Karras et al. (2022)."""
    ramp = th.linspace(0, 1, n)
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    return append_zero(sigmas).to(device)


def to_d(x, sigma, denoised):
    """Converts a denoiser output to a Karras ODE derivative."""
    return (x - denoised) / append_dims(sigma, x.ndim)


def get_ancestral_step(sigma_from, sigma_to):
    """Calculates the noise level (sigma_down) to step down to and the amount
    of noise to add (sigma_up) when doing an ancestral sampling step."""
    sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
    sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
    return sigma_down, sigma_up


@th.no_grad()
def sample_euler_ancestral(model, x, sigmas, generator, progress=False, callback=None):
    """Ancestral sampling with Euler method steps."""
    s_in = x.new_ones([x.shape[0]])
    indices = range(len(sigmas) - 1)
    if progress:
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        denoised = model(x, sigmas[i] * s_in)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
        if callback is not None:
            callback(
                {
                    "x": x,
                    "i": i,
                    "sigma": sigmas[i],
                    "sigma_hat": sigmas[i],
                    "denoised": denoised,
                }
            )
        d = to_d(x, sigmas[i], denoised)
        # Euler method
        dt = sigma_down - sigmas[i]
        x = x + d * dt
        x = x + generator.randn_like(x) * sigma_up
    return x


@th.no_grad()
def sample_midpoint_ancestral(model, x, ts, generator, progress=False, callback=None):
    """Ancestral sampling with midpoint method steps."""
    s_in = x.new_ones([x.shape[0]])
    step_size = 1 / len(ts)
    if progress:
        from tqdm.auto import tqdm

        ts = tqdm(ts)

    for tn in ts:
        dn = model(x, tn * s_in)
        dn_2 = model(x + (step_size / 2) * dn, (tn + step_size / 2) * s_in)
        x = x + step_size * dn_2
        if callback is not None:
            callback({"x": x, "tn": tn, "dn": dn, "dn_2": dn_2})
    return x


@th.no_grad()
def sample_heun(
    denoiser,
    x,
    sigmas,
    generator,
    progress=False,
    callback=None,
    s_churn=0.0,
    s_tmin=0.0,
    s_tmax=float("inf"),
    s_noise=1.0,
):
    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
    s_in = x.new_ones([x.shape[0]])
    indices = range(len(sigmas) - 1)
    if progress:
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        gamma = (
            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1)
            if s_tmin <= sigmas[i] <= s_tmax
            else 0.0
        )
        eps = generator.randn_like(x) * s_noise
        sigma_hat = sigmas[i] * (gamma + 1)
        if gamma > 0:
            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
        denoised = denoiser(x, sigma_hat * s_in)
        d = to_d(x, sigma_hat, denoised)
        if callback is not None:
            callback(
                {
                    "x": x,
                    "i": i,
                    "sigma": sigmas[i],
                    "sigma_hat": sigma_hat,
                    "denoised": denoised,
                }
            )
        dt = sigmas[i + 1] - sigma_hat
        if sigmas[i + 1] == 0:
            # Euler method
            x = x + d * dt
        else:
            # Heun's method
            x_2 = x + d * dt
            denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
            d_prime = (d + d_2) / 2
            x = x + d_prime * dt
    return x


@th.no_grad()
def sample_euler(
    denoiser,
    x,
    sigmas,
    generator,
    progress=False,
    callback=None,
):
    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
    s_in = x.new_ones([x.shape[0]])
    indices = range(len(sigmas) - 1)
    if progress:
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        sigma = sigmas[i]
        denoised = denoiser(x, sigma * s_in)
        d = to_d(x, sigma, denoised)
        if callback is not None:
            callback(
                {
                    "x": x,
                    "i": i,
                    "sigma": sigmas[i],
                    "denoised": denoised,
                }
            )
        dt = sigmas[i + 1] - sigma
        x = x + d * dt
    return x


@th.no_grad()
def sample_dpm(
    denoiser,
    x,
    sigmas,
    generator,
    progress=False,
    callback=None,
    s_churn=0.0,
    s_tmin=0.0,
    s_tmax=float("inf"),
    s_noise=1.0,
):
    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
    s_in = x.new_ones([x.shape[0]])
    indices = range(len(sigmas) - 1)
    if progress:
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        gamma = (
            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1)
            if s_tmin <= sigmas[i] <= s_tmax
            else 0.0
        )
        eps = generator.randn_like(x) * s_noise
        sigma_hat = sigmas[i] * (gamma + 1)
        if gamma > 0:
            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
        denoised = denoiser(x, sigma_hat * s_in)
        d = to_d(x, sigma_hat, denoised)
        if callback is not None:
            callback(
                {
                    "x": x,
                    "i": i,
                    "sigma": sigmas[i],
                    "sigma_hat": sigma_hat,
                    "denoised": denoised,
                }
            )
        # Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
        sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2) ** 3
        dt_1 = sigma_mid - sigma_hat
        dt_2 = sigmas[i + 1] - sigma_hat
        x_2 = x + d * dt_1
        denoised_2 = denoiser(x_2, sigma_mid * s_in)
        d_2 = to_d(x_2, sigma_mid, denoised_2)
        x = x + d_2 * dt_2
    return x


@th.no_grad()
def sample_onestep(
    distiller,
    x,
    sigmas,
    generator=None,
    progress=False,
    callback=None,
):
    """Single-step generation from a distilled model."""
    s_in = x.new_ones([x.shape[0]])
    return distiller(x, sigmas[0] * s_in)


@th.no_grad()
def stochastic_iterative_sampler(
    distiller,
    x,
    sigmas,
    generator,
    ts,
    progress=False,
    callback=None,
    t_min=0.002,
    t_max=80.0,
    rho=7.0,
    steps=40,
):
    t_max_rho = t_max ** (1 / rho)
    t_min_rho = t_min ** (1 / rho)
    s_in = x.new_ones([x.shape[0]])

    for i in range(len(ts) - 1):
        t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
        x0 = distiller(x, t * s_in)
        next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
        next_t = np.clip(next_t, t_min, t_max)
        x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)

    return x


@th.no_grad()
def sample_progdist(
    denoiser,
    x,
    sigmas,
    generator=None,
    progress=False,
    callback=None,
):
    s_in = x.new_ones([x.shape[0]])
    sigmas = sigmas[:-1]  # skip the zero sigma

    indices = range(len(sigmas) - 1)
    if progress:
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        sigma = sigmas[i]
        denoised = denoiser(x, sigma * s_in)
        d = to_d(x, sigma, denoised)
        if callback is not None:
            callback(
                {
                    "x": x,
                    "i": i,
                    "sigma": sigma,
                    "denoised": denoised,
                }
            )
        dt = sigmas[i + 1] - sigma
        x = x + d * dt

    return x


# @th.no_grad()
# def iterative_colorization(
#     distiller,
#     images,
#     x,
#     ts,
#     t_min=0.002,
#     t_max=80.0,
#     rho=7.0,
#     steps=40,
#     generator=None,
# ):
#     def obtain_orthogonal_matrix():
#         vector = np.asarray([0.2989, 0.5870, 0.1140])
#         vector = vector / np.linalg.norm(vector)
#         matrix = np.eye(3)
#         matrix[:, 0] = vector
#         matrix = np.linalg.qr(matrix)[0]
#         if np.sum(matrix[:, 0]) < 0:
#             matrix = -matrix
#         return matrix

#     Q = th.from_numpy(obtain_orthogonal_matrix()).to(dist_util.dev()).to(th.float32)
#     mask = th.zeros(*x.shape[1:], device=dist_util.dev())
#     mask[0, ...] = 1.0

#     def replacement(x0, x1):
#         x0 = th.einsum("bchw,cd->bdhw", x0, Q)
#         x1 = th.einsum("bchw,cd->bdhw", x1, Q)

#         x_mix = x0 * mask + x1 * (1.0 - mask)
#         x_mix = th.einsum("bdhw,cd->bchw", x_mix, Q)
#         return x_mix

#     t_max_rho = t_max ** (1 / rho)
#     t_min_rho = t_min ** (1 / rho)
#     s_in = x.new_ones([x.shape[0]])
#     images = replacement(images, th.zeros_like(images))

#     for i in range(len(ts) - 1):
#         t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         x0 = distiller(x, t * s_in)
#         x0 = th.clamp(x0, -1.0, 1.0)
#         x0 = replacement(images, x0)
#         next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         next_t = np.clip(next_t, t_min, t_max)
#         x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)

#     return x, images


# @th.no_grad()
# def iterative_inpainting(
#     distiller,
#     images,
#     x,
#     ts,
#     t_min=0.002,
#     t_max=80.0,
#     rho=7.0,
#     steps=40,
#     generator=None,
# ):
#     from PIL import Image, ImageDraw, ImageFont

#     image_size = x.shape[-1]

#     # create a blank image with a white background
#     img = Image.new("RGB", (image_size, image_size), color="white")

#     # get a drawing context for the image
#     draw = ImageDraw.Draw(img)

#     # load a font
#     font = ImageFont.truetype("arial.ttf", 250)

#     # draw the letter "C" in black
#     draw.text((50, 0), "S", font=font, fill=(0, 0, 0))

#     # convert the image to a numpy array
#     img_np = np.array(img)
#     img_np = img_np.transpose(2, 0, 1)
#     img_th = th.from_numpy(img_np).to(dist_util.dev())

#     mask = th.zeros(*x.shape, device=dist_util.dev())
#     mask = mask.reshape(-1, 7, 3, image_size, image_size)

#     mask[::2, :, img_th > 0.5] = 1.0
#     mask[1::2, :, img_th < 0.5] = 1.0
#     mask = mask.reshape(-1, 3, image_size, image_size)

#     def replacement(x0, x1):
#         x_mix = x0 * mask + x1 * (1 - mask)
#         return x_mix

#     t_max_rho = t_max ** (1 / rho)
#     t_min_rho = t_min ** (1 / rho)
#     s_in = x.new_ones([x.shape[0]])
#     images = replacement(images, -th.ones_like(images))

#     for i in range(len(ts) - 1):
#         t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         x0 = distiller(x, t * s_in)
#         x0 = th.clamp(x0, -1.0, 1.0)
#         x0 = replacement(images, x0)
#         next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         next_t = np.clip(next_t, t_min, t_max)
#         x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)

#     return x, images


# @th.no_grad()
# def iterative_superres(
#     distiller,
#     images,
#     x,
#     ts,
#     t_min=0.002,
#     t_max=80.0,
#     rho=7.0,
#     steps=40,
#     generator=None,
# ):
#     patch_size = 8

#     def obtain_orthogonal_matrix():
#         vector = np.asarray([1] * patch_size**2)
#         vector = vector / np.linalg.norm(vector)
#         matrix = np.eye(patch_size**2)
#         matrix[:, 0] = vector
#         matrix = np.linalg.qr(matrix)[0]
#         if np.sum(matrix[:, 0]) < 0:
#             matrix = -matrix
#         return matrix

#     Q = th.from_numpy(obtain_orthogonal_matrix()).to(dist_util.dev()).to(th.float32)

#     image_size = x.shape[-1]

#     def replacement(x0, x1):
#         x0_flatten = (
#             x0.reshape(-1, 3, image_size, image_size)
#             .reshape(
#                 -1,
#                 3,
#                 image_size // patch_size,
#                 patch_size,
#                 image_size // patch_size,
#                 patch_size,
#             )
#             .permute(0, 1, 2, 4, 3, 5)
#             .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
#         )
#         x1_flatten = (
#             x1.reshape(-1, 3, image_size, image_size)
#             .reshape(
#                 -1,
#                 3,
#                 image_size // patch_size,
#                 patch_size,
#                 image_size // patch_size,
#                 patch_size,
#             )
#             .permute(0, 1, 2, 4, 3, 5)
#             .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
#         )
#         x0 = th.einsum("bcnd,de->bcne", x0_flatten, Q)
#         x1 = th.einsum("bcnd,de->bcne", x1_flatten, Q)
#         x_mix = x0.new_zeros(x0.shape)
#         x_mix[..., 0] = x0[..., 0]
#         x_mix[..., 1:] = x1[..., 1:]
#         x_mix = th.einsum("bcne,de->bcnd", x_mix, Q)
#         x_mix = (
#             x_mix.reshape(
#                 -1,
#                 3,
#                 image_size // patch_size,
#                 image_size // patch_size,
#                 patch_size,
#                 patch_size,
#             )
#             .permute(0, 1, 2, 4, 3, 5)
#             .reshape(-1, 3, image_size, image_size)
#         )
#         return x_mix

#     def average_image_patches(x):
#         x_flatten = (
#             x.reshape(-1, 3, image_size, image_size)
#             .reshape(
#                 -1,
#                 3,
#                 image_size // patch_size,
#                 patch_size,
#                 image_size // patch_size,
#                 patch_size,
#             )
#             .permute(0, 1, 2, 4, 3, 5)
#             .reshape(-1, 3, image_size**2 // patch_size**2, patch_size**2)
#         )
#         x_flatten[..., :] = x_flatten.mean(dim=-1, keepdim=True)
#         return (
#             x_flatten.reshape(
#                 -1,
#                 3,
#                 image_size // patch_size,
#                 image_size // patch_size,
#                 patch_size,
#                 patch_size,
#             )
#             .permute(0, 1, 2, 4, 3, 5)
#             .reshape(-1, 3, image_size, image_size)
#         )

#     t_max_rho = t_max ** (1 / rho)
#     t_min_rho = t_min ** (1 / rho)
#     s_in = x.new_ones([x.shape[0]])
#     images = average_image_patches(images)

#     for i in range(len(ts) - 1):
#         t = (t_max_rho + ts[i] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         x0 = distiller(x, t * s_in)
#         x0 = th.clamp(x0, -1.0, 1.0)
#         x0 = replacement(images, x0)
#         next_t = (t_max_rho + ts[i + 1] / (steps - 1) * (t_min_rho - t_max_rho)) ** rho
#         next_t = np.clip(next_t, t_min, t_max)
#         x = x0 + generator.randn_like(x) * np.sqrt(next_t**2 - t_min**2)

#     return x, images