sample.py

#!/usr/bin/env python3

# Code by Kat Crowson in k-diffusion repo, modified by Scott H Hawley (SHH)
# Modified by Scott H. Hawley for masking, ZeroGPU ets.

"""Samples from k-diffusion models."""

import gradio 
import spaces
import natten
import argparse
from pathlib import Path

import accelerate
import safetensors.torch as safetorch
import torch
from tqdm import trange, tqdm
from PIL import Image
from torchvision import transforms

import k_diffusion as K

from pom.v_diffusion import DDPM, LogSchedule, CrashSchedule
#CHORD_BORDER = 8   # chord border size in pixels
from pom.chords import CHORD_BORDER, img_batch_to_seq_emb, ChordSeqEncoder


# ---- my mangled sampler that includes repaint 
import torchsde 

@spaces.GPU
class BatchedBrownianTree:
    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""

    def __init__(self, x, t0, t1, seed=None, **kwargs):
        t0, t1, self.sign = self.sort(t0, t1)
        w0 = kwargs.get('w0', torch.zeros_like(x))
        if seed is None:
            seed = torch.randint(0, 2 ** 63 - 1, []).item()
        self.batched = True
        try:
            assert len(seed) == x.shape[0]
            w0 = w0[0]
        except TypeError:
            seed = [seed]
            self.batched = False
        self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]

    @staticmethod
    def sort(a, b):
        return (a, b, 1) if a < b else (b, a, -1)

    def __call__(self, t0, t1):
        t0, t1, sign = self.sort(t0, t1)
        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
        return w if self.batched else w[0]


@spaces.GPU
class BrownianTreeNoiseSampler:
    """A noise sampler backed by a torchsde.BrownianTree.

    Args:
        x (Tensor): The tensor whose shape, device and dtype to use to generate
            random samples.
        sigma_min (float): The low end of the valid interval.
        sigma_max (float): The high end of the valid interval.
        seed (int or List[int]): The random seed. If a list of seeds is
            supplied instead of a single integer, then the noise sampler will
            use one BrownianTree per batch item, each with its own seed.
        transform (callable): A function that maps sigma to the sampler's
            internal timestep.
    """

    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
        self.transform = transform
        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
        self.tree = BatchedBrownianTree(x, t0, t1, seed)

    def __call__(self, sigma, sigma_next):
        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()

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 to_d(x, sigma, denoised):
    """Converts a denoiser output to a Karras ODE derivative."""
    return (x - denoised) / append_dims(sigma, x.ndim)


@spaces.GPU
@torch.no_grad()
def my_sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., repaint=1):
    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    for i in trange(len(sigmas) - 1, disable=disable):
        for u in range(repaint):
            gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
            eps = torch.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 = model(x, sigma_hat * s_in, **extra_args)
            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
            # Euler method
            x = x + d * dt
            if x.isnan().any():
                assert False, f"x has NaNs, i = {i}, u = {u}, repaint = {repaint}"
            if u < repaint - 1:
                beta = (sigmas[i + 1] / sigmas[-1]) ** 2
                x = torch.sqrt(1 - beta) * x + torch.sqrt(beta) * torch.randn_like(x)

    return x

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


@spaces.GPU
@torch.no_grad()
def my_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, 
                    disable=None, eta=1., s_noise=1., noise_sampler=None, 
                    solver_type='midpoint',
                    repaint=4):
    """DPM-Solver++(2M) SDE.  but with repaint added"""

    if solver_type not in {'heun', 'midpoint'}:
        raise ValueError('solver_type must be \'heun\' or \'midpoint\'')

    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    old_denoised = None
    h_last = None
    old_x = None

    for i in trange(len(sigmas) - 1, disable=disable):  # time loop

        for u in range(repaint):
            denoised = model(x, sigmas[i] * s_in, **extra_args)
            if callback is not None:
                callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
            #print("i, u, sigmas[i], sigmas[i + 1] = ", i, u, sigmas[i], sigmas[i + 1])
            if sigmas[i + 1] == 0:
                # Denoising step
                x = denoised
            else:
                # DPM-Solver++(2M) SDE
                t, s = -sigmas[i].log(), -sigmas[i + 1].log()
                h = s - t
                eta_h = eta * h

                x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised

                if old_denoised is not None:
                    r = h_last / h
                    if solver_type == 'heun':
                        x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
                    elif solver_type == 'midpoint':
                        x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)

                if eta:
                    x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
                
                
                if callback is not None:
                    callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})   

                if x.isnan().any():
                    assert False, f"x has NaNs, i = {i}, u = {u}, repaint = {repaint}"
            
                if u < repaint - 1:
                    # RePaint: go "back" in integration via the "forward" process, by adding a little noise to x
                    #  ...but scaled properly!
                    # But how to convert from original RePaint to k-diffusion?  I'll try a few variants
                    repaint_choice = 'orig' # ['orig','var1','var2', etc...]

                    sigma_diff = (sigmas[i] - sigmas[i+1]).abs()
                    sigma_ratio = ( sigmas[i+1] / sigma_max ) # use i+1 or i?
                    if repaint_choice == 'orig': # attempt at original RePaint algorithm, which used betas
                        # if sigmas are the std devs, then betas are variances?  but beta_max = 1, so how to get that? ratio?
                        beta = sigma_ratio**2
                        x = torch.sqrt(1-beta)*x +  torch.sqrt(beta)*torch.randn_like(x) # this is from RePaint Paper
                    elif repaint_choice == 'var1': # or maybe this...?  # worse than orig
                        x = x + sigma_diff*torch.randn_like(x)
                    elif repaint_choice == 'var2':  # or this...?  # yields NaNs
                        x = (1-sigma_diff)*x + sigma_diff*torch.randn_like(x)
                    elif repaint_choice == 'var3':            # results similar to var1
                        x = (1.0-sigma_ratio)*x + sigmas[i+1]*torch.randn_like(x)
                    elif repaint_choice == 'var4':   # NaNs         # stealing code from elsewhere, no idea WTF I'm doing.
                        #Invert this: target = (input - c_skip * noised_input) / c_out, where target = model_output
                        x_tm1, x_t = x, old_x 
                        #              x_tm1 = ( x_0  - c_skip * noised_x0 ) / c_out
                        #       So     x_tm1*c_out = x_0 - c_skip * noised_x0
                        input, noise = x_tm1, torch.randn_like(x)
                        noised_input = input + noise * append_dims(sigma_diff, input.ndim)
                        c_skip, c_out, c_in = [append_dims(x, input.ndim) for x in get_scalings(sigmas[i])]
                        model_output = x_tm1
                        renoised_x = c_out * model_output + c_skip * noised_input 
                        x = renoised_x
                    elif repaint_choice == 'var5':
                        x = torch.sqrt((1-(sigma_diff/sigma_max)**2))*x + sigma_diff*torch.randn_like(x)

                    # include this?  guessing no.
                    #old_denoised = denoised
                    #h_last = h

        old_denoised = denoised
        h_last = h
        old_x = x
    return x




# -----from stable-audio-tools

# Define the noise schedule and sampling loop
def get_alphas_sigmas(t):
    """Returns the scaling factors for the clean image (alpha) and for the
    noise (sigma), given a timestep."""
    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)

def alpha_sigma_to_t(alpha, sigma):
    """Returns a timestep, given the scaling factors for the clean image and for
    the noise."""
    return torch.atan2(sigma, alpha) / math.pi * 2

def t_to_alpha_sigma(t):
    """Returns the scaling factors for the clean image and for the noise, given
    a timestep."""
    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)

@torch.no_grad()
def sample(model, x, steps, eta, **extra_args):
    """Draws samples from a model given starting noise. v-diffusion"""
    ts = x.new_ones([x.shape[0]])

    # Create the noise schedule
    t = torch.linspace(1, 0, steps + 1)[:-1]

    alphas, sigmas = get_alphas_sigmas(t)

    # The sampling loop
    for i in trange(steps):

        # Get the model output (v, the predicted velocity)
        with torch.cuda.amp.autocast():
            v = model(x, ts * t[i], **extra_args).float()

        # Predict the noise and the denoised image
        pred = x * alphas[i] - v * sigmas[i]
        eps = x * sigmas[i] + v * alphas[i]

        # If we are not on the last timestep, compute the noisy image for the
        # next timestep.
        if i < steps - 1:
            # If eta > 0, adjust the scaling factor for the predicted noise
            # downward according to the amount of additional noise to add
            ddim_sigma = eta * (sigmas[i + 1]**2 / sigmas[i]**2).sqrt() * \
                (1 - alphas[i]**2 / alphas[i + 1]**2).sqrt()
            adjusted_sigma = (sigmas[i + 1]**2 - ddim_sigma**2).sqrt()

            # Recombine the predicted noise and predicted denoised image in the
            # correct proportions for the next step
            x = pred * alphas[i + 1] + eps * adjusted_sigma

            # Add the correct amount of fresh noise
            if eta:
                x += torch.randn_like(x) * ddim_sigma

    # If we are on the last timestep, output the denoised image
    return pred

# Soft mask inpainting is just shrinking hard (binary) mask inpainting
# Given a float-valued soft mask (values between 0 and 1), get the binary mask for this particular step
@spaces.GPU
def get_bmask(i, steps, mask):
    strength = (i+1)/(steps)
    # convert to binary mask
    bmask = torch.where(mask<=strength,1,0)
    return bmask

@spaces.GPU
def make_cond_model_fn(model, cond_fn):
    def cond_model_fn(x, sigma, **kwargs):
        with torch.enable_grad():
            x = x.detach().requires_grad_()
            denoised = model(x, sigma, **kwargs)
            cond_grad = cond_fn(x, sigma, denoised=denoised, **kwargs).detach()
            cond_denoised = denoised.detach() + cond_grad * K.utils.append_dims(sigma**2, x.ndim)
        return cond_denoised
    return cond_model_fn

# Uses k-diffusion from https://github.com/crowsonkb/k-diffusion
# init_data is init_audio as latents (if this is latent diffusion)
# For sampling, set both init_data and mask to None
# For variations, set init_data 
# For inpainting, set both init_data & mask 
@spaces.GPU
def sample_k(
        model_fn, 
        noise, 
        init_data=None,
        mask=None,
        steps=100, 
        sampler_type="dpmpp-2m-sde", 
        sigma_min=0.5, 
        sigma_max=50, 
        rho=1.0, device="cuda", 
        callback=None, 
        cond_fn=None,
        model_config=None,
        repaint=1,
        **extra_args
    ):

    #denoiser = K.external.VDenoiser(model_fn)
    denoiser = K.Denoiser(model_fn, sigma_data=model_config['sigma_data'])

    if cond_fn is not None:
        denoiser = make_cond_model_fn(denoiser, cond_fn)

    # Make the list of sigmas. Sigma values are scalars related to the amount of noise each denoising step has
    #sigmas = K.sampling.get_sigmas_polyexponential(steps, sigma_min, sigma_max, rho, device=device)
    sigmas = K.sampling.get_sigmas_karras(steps, sigma_min, sigma_max, rho=7., device=device)
    print("sigmas[0] = ", sigmas[0])
    # Scale the initial noise by sigma 
    noise = noise * sigmas[0]

    wrapped_callback = callback

    if mask is None and init_data is not None:
        # VARIATION (no inpainting)
        # set the initial latent to the init_data, and noise it with initial sigma
        x = init_data + noise 
    elif mask is not None and init_data is not None:
        # INPAINTING
        bmask = get_bmask(0, steps, mask)
        # initial noising
        input_noised = init_data + noise
        # set the initial latent to a mix of init_data and noise, based on step 0's binary mask
        x = input_noised * bmask + noise * (1-bmask)
        # define the inpainting callback function (Note: side effects, it mutates x)
        # See https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L596C13-L596C105
        # callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        # This is called immediately after `denoised = model(x, sigmas[i] * s_in, **extra_args)`
        def inpainting_callback(args):
            i = args["i"]
            x = args["x"]
            sigma = args["sigma"]
            #denoised = args["denoised"]
            # noise the init_data input with this step's appropriate amount of noise
            input_noised = init_data + torch.randn_like(init_data) * sigma
            # shrinking hard mask
            bmask = get_bmask(i, steps, mask)
            # mix input_noise with x, using binary mask
            new_x = input_noised * bmask + x * (1-bmask)
            # mutate x
            x[:,:,:] = new_x[:,:,:]
        # wrap together the inpainting callback and the user-submitted callback. 
        if callback is None: 
            wrapped_callback = inpainting_callback
        else:
            wrapped_callback = lambda args: (inpainting_callback(args), callback(args))
    else:
        # SAMPLING
        # set the initial latent to noise
        x = noise


    print("sample_k: x.min, x.max = ", x.min(), x.max())
    print(f"sample_k: key, val.dtype = ",[ (key, val.dtype if val is not None else val) for key,val in extra_args.items()])
    with torch.cuda.amp.autocast():
        if sampler_type == "k-heun":
            return K.sampling.sample_heun(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "k-lms":
            return K.sampling.sample_lms(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "k-dpmpp-2s-ancestral":
            return K.sampling.sample_dpmpp_2s_ancestral(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "k-dpm-2":
            return K.sampling.sample_dpm_2(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "k-dpm-fast":
            return K.sampling.sample_dpm_fast(denoiser, x, sigma_min, sigma_max, steps, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "k-dpm-adaptive":
            return K.sampling.sample_dpm_adaptive(denoiser, x, sigma_min, sigma_max, rtol=0.01, atol=0.01, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "dpmpp-2m-sde":
            return K.sampling.sample_dpmpp_2m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "my-dpmpp-2m-sde":
            return my_dpmpp_2m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, repaint=repaint, extra_args=extra_args)
        elif sampler_type == "dpmpp-3m-sde":
            return K.sampling.sample_dpmpp_3m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
        elif sampler_type == "my-sample-euler":
            return my_sample_euler(denoiser, x, sigmas, disable=False, callback=wrapped_callback, repaint=repaint, extra_args=extra_args)


## ---- end stable-audio-tools
@spaces.GPU
def infer_mask_from_init_img(img, mask_with='white'):
    """given an image with mask areas marked, extract the mask itself
       note, this works whether image is normalized on 0..1 or -1..1, but not 0..255"""
    print("Inferring mask from init_img")
    assert mask_with in ['blue','white']
    if not torch.is_tensor(img):
        img = ToTensor()(img)
    mask = torch.zeros(img.shape[-2:])
    if mask_with == 'white':
        mask[ (img[0,:,:]==1) & (img[1,:,:]==1) & (img[2,:,:]==1)] = 1
    elif mask_with == 'blue':
        mask[img[2,:,:]==1] = 1  # blue
    return mask*1.0

@spaces.GPU
def grow_mask(init_mask, grow_by=2):
    "adds a border of grow_by pixels to the mask, by growing it grow_by times. If grow_by=0, does nothing"
    new_mask = init_mask.clone()
    for c in range(grow_by):
        # wherever mask is bordered by a 1, set it to 1
        new_mask[1:-1,1:-1] = (new_mask[1:-1,1:-1] + new_mask[0:-2,1:-1] + new_mask[2:,1:-1] + new_mask[1:-1,0:-2] + new_mask[1:-1,2:]) > 0 
    return new_mask

@spaces.GPU
def add_seeding(init_image, init_mask, grow_by=0, seed_scale=1.0):
    "adds extra noise inside mask"
    init_mask = grow_mask(init_mask, grow_by=grow_by)  # make the mask bigger
    if not torch.is_tensor(init_image):
        init_image = ToTensor()(init_image)
    init_image = init_image.clone()
    # wherever mask is 1, set first set init_image to min value 
    init_image[:,init_mask == 1] = init_image.min()   
    init_image = init_image + seed_scale*torch.randn_like(init_image) * (init_mask) # add noise where mask is 1
    # wherever the mask is 1, set the blue channel to -1.0, otherwise leave it alone
    init_image[2,:,:] = init_image[2,:,:] * (1-init_mask) - 1.0*init_mask
    return init_image

@spaces.GPU
def get_init_image_and_mask(args, device):
    convert_tensor = transforms.ToTensor()
    init_image = Image.open(args.init_image).convert('RGB')
    init_image = convert_tensor(init_image)
    #normalize image from 0..1 to -1..1
    init_image = (2.0 * init_image) - 1.0


    init_mask = torch.ones(init_image.shape[-2:])  # ones are where stuff will change, zeros will stay the same

    inpaint_task = 'infer'  # infer mask from init_image
    assert inpaint_task in ['accomp','chords','melody','nucleation','notes','continue','infer']

    if inpaint_task in ['melody','accomp']:
        init_mask[0:70,:] = 0 # zero out a melody strip of image near top
        init_mask[128+0:128+70,:] = 0 # zero out a melody strip of image along bottom row
        if inpaint_task == 'melody':
            init_mask = 1 - init_mask 
    elif inpaint_task in ['notes','chords']:
        # keep chords only
        #init_mask = torch.ones_like(x) 
        init_mask[0:CHORD_BORDER,:] = 0  # top row of 256x256
        init_mask[128-CHORD_BORDER:128+CHORD_BORDER,:] = 0  # middle rows of 256x256
        init_mask[-CHORD_BORDER:,:] = 0  # bottom row of 256x256
        if inpaint_task == 'chords':
            init_mask = 1 - init_mask # inverse: genereate chords given notes
    elif inpaint_task == 'continue': 
        init_mask[0:128,:] = 0     # remember it's a square, so just mask out the bottom half
    elif inpaint_task == 'nucleation':
        # set mask to wherever the blue channel is >= 0.9
        init_mask = (init_image[2,:,:] > 0.0)*1.0
        # zero out init mask in top and bottom borders
        init_mask[0:CHORD_BORDER,:] = 0
        init_mask[-CHORD_BORDER:,:] = 0
        init_mask[128-CHORD_BORDER:128+CHORD_BORDER,:] = 0

        # remove all blue in init_image between the borders
        init_image[2,CHORD_BORDER:128-CHORD_BORDER,:] = -1.0
        init_image[2,128+CHORD_BORDER:-CHORD_BORDER,:] = -1.0

        # grow the sides of the mask by one pixel:
        # wherever mask is zero but is bordered by a 1, set it to 1
        init_mask[1:-1,1:-1] = (init_mask[1:-1,1:-1] + init_mask[0:-2,1:-1] + init_mask[2:,1:-1] + init_mask[1:-1,0:-2] + init_mask[1:-1,2:]) > 0 
        #init_mask[1:-1,1:-1] = (init_mask[1:-1,1:-1] + init_mask[0:-2,1:-1] + init_mask[2:,1:-1] + init_mask[1:-1,0:-2] + init_mask[1:-1,2:]) > 0 
    elif inpaint_task == 'infer':
        init_mask = infer_mask_from_init_img(init_image, mask_with='white')

    # Also black out init_image wherever init mask is 1 
    init_image[:,init_mask == 1] = init_image.min()

    if args.seed_scale > 0: # driving nucleation
        print("Seeding nucleation, seed_scale = ", args.seed_scale)
        init_image = add_seeding(init_image, init_mask, grow_by=0, seed_scale=args.seed_scale)

    # remove any blue in middle of init image
    print("init_image.shape = ", init_image.shape)
    init_image[2,CHORD_BORDER:128-CHORD_BORDER,:] = -1.0
    init_image[2,128+CHORD_BORDER:-CHORD_BORDER,:] = -1.0

    # Debugging: output some images so we can see what's going on
    init_mask_t = init_mask.float()*255 # convert mask to 0..255 for writing as image
    # Convert to NumPy array and rearrange dimensions
    init_mask_img_numpy = init_mask_t.byte().cpu().numpy()#.transpose(1, 2, 0)
    init_mask_debug_img = Image.fromarray(init_mask_img_numpy)
    init_mask_debug_img.save("init_mask_debug.png")
    init_image_debug_img = Image.fromarray((init_image*127.5+127.5).byte().cpu().numpy().transpose(1,2,0))
    init_image_debug_img.save("init_image_debug.png")

    # reshape image and mask to be 4D tensors
    init_image = init_image.unsqueeze(0).repeat(args.batch_size, 1, 1, 1)
    init_mask = init_mask.unsqueeze(0).unsqueeze(1).repeat(args.batch_size,3,1,1).float()
    return init_image.to(device), init_mask.to(device)

#@spaces.GPU # generates an error
def main():
    global init_image, init_mask
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument('--batch-size', type=int, default=64,
                   help='the batch size')
    p.add_argument('--checkpoint', type=Path, required=True,
                   help='the checkpoint to use')
    p.add_argument('--config', type=Path,
                   help='the model config')
    p.add_argument('-n', type=int, default=64,
                   help='the number of images to sample')
    p.add_argument('--prefix', type=str, default='out',
                   help='the output prefix')
    p.add_argument('--repaint', type=int, default=1,
                   help='number of (re)paint steps')
    p.add_argument('--steps', type=int, default=50,
                   help='the number of denoising steps')
    p.add_argument('--seed-scale', type=float, default=0.0, help='strength of nucleation seeding')
    p.add_argument('--init-image', type=Path, default=None, help='the initial image')
    p.add_argument('--init-strength', type=float, default=1., help='strength of init image')
    args = p.parse_args()
    print("args =", args, flush=True)

    config = K.config.load_config(args.config if args.config else args.checkpoint)
    model_config = config['model']
    # TODO: allow non-square input sizes
    assert len(model_config['input_size']) == 2 and model_config['input_size'][0] == model_config['input_size'][1]
    size = model_config['input_size']

    accelerator = accelerate.Accelerator()
    device = accelerator.device
    print('Using device:', device, flush=True)

    inner_model = K.config.make_model(config).eval().requires_grad_(False).to(device)
    cse = None # ChordSeqEncoder().eval().requires_grad_(False).to(device)  # add chord embedding-maker to main model
    if cse is not None:
        inner_model.cse = cse
    try:
        inner_model.load_state_dict(safetorch.load_file(args.checkpoint))
    except:
        #ckpt = torch.load(args.checkpoint).to(device)
        ckpt = torch.load(args.checkpoint, map_location='cpu')
        inner_model.load_state_dict(ckpt['model'])

    accelerator.print('Parameters:', K.utils.n_params(inner_model))
    model = K.Denoiser(inner_model, sigma_data=model_config['sigma_data'])

    sigma_min = model_config['sigma_min']
    sigma_max = model_config['sigma_max']

    # SHH modified
    torch.set_float32_matmul_precision('high')
    #class_cond = torch.tensor([0]).to(device)
    #num_classes = 10
    #class_cond = torch.remainder(torch.arange(0, args.n), num_classes).int().to(device)
    #extra_args = {'class_cond':class_cond}
    extra_args = {}
    init_image, init_mask = None, None
    if args.init_image is not None:
        init_image, init_mask = get_init_image_and_mask(args, device)
        init_image = init_image.to(device)
        init_mask = init_mask.to(device)

    @torch.no_grad()
    @K.utils.eval_mode(model)
    def run():
        global init_image, init_mask
        if accelerator.is_local_main_process:
            tqdm.write('Sampling...')
        sigmas = K.sampling.get_sigmas_karras(args.steps, sigma_min, sigma_max, rho=7., device=device)

        #ddpm_sampler = DDPM(model)
        #model_fn = model
        #ddpm_sampler = K.external.VDenoiser(model_fn)

        @spaces.GPU
        def sample_fn(n, debug=True):
            x = torch.randn([n, model_config['input_channels'], size[0], size[1]], device=device) * sigma_max
            print("n, sigma_max, x.min, x.max = ", n, sigma_max, x.min(), x.max())

            if args.init_image is not None:
                init_data, mask = get_init_image_and_mask(args, device)
                init_data = args.seed_scale*x*mask + (1-mask)*init_data  # extra nucleation?
                if cse is not None: 
                    chord_cond = img_batch_to_seq_emb(init_data, inner_model.cse).to(device)
                else: 
                    chord_cond = None
                #print("init_data.shape, init_data.min, init_data.max = ", init_data.shape, init_data.min(), init_data.max())
            else:
                init_data, mask, chord_cond = None, None, None
            # chord_cond doesn't work anyway so f it: 
            chord_cond = None

            print("chord_cond = ", chord_cond)
            if chord_cond is not None: 
                extra_args['chord_cond'] = chord_cond
            # these two work:
            #x_0 = K.sampling.sample_lms(model, x, sigmas, disable=not accelerator.is_local_main_process, extra_args=extra_args)
            #x_0 = K.sampling.sample_dpmpp_2m_sde(model, x, sigmas, disable=not accelerator.is_local_main_process, extra_args=extra_args)

            noise = torch.randn([n, model_config['input_channels'], size[0], size[1]], device=device) 

            sampler_type="my-dpmpp-2m-sde"  # "k-lms"
            #sampler_type="my-sample-euler"
            #sampler_type="dpmpp-2m-sde"  
            #sampler_type = "dpmpp-3m-sde"
            #sampler_type = "k-dpmpp-2s-ancestral"
            print("dtypes:", [x.dtype if x is not None else None  for x in [noise, init_data, mask, chord_cond]])
            x_0 = sample_k(inner_model, noise, sampler_type=sampler_type, 
                           init_data=init_data, mask=mask, steps=args.steps, 
                           sigma_min=sigma_min, sigma_max=sigma_max, rho=7., 
                           device=device, model_config=model_config, repaint=args.repaint, 
                           **extra_args)
            #x_0 = sample_k(inner_model, noise, sampler_type="dpmpp-2m-sde", steps=100,  sigma_min=0.5, sigma_max=50, rho=1., device=device,  model_config=model_config, **extra_args)
            print("x_0.min, x_0.max = ", x_0.min(), x_0.max())
            if x_0.isnan().any():
                assert False, "x_0 has NaNs"
            
            # do gpu garbage collection before proceeding
            torch.cuda.empty_cache()
            return x_0
        
        x_0 = K.evaluation.compute_features(accelerator, sample_fn, lambda x: x, args.n, args.batch_size)
        if accelerator.is_main_process:
            for i, out in enumerate(x_0):
                filename = f'{args.prefix}_{i:05}.png'
                K.utils.to_pil_image(out).save(filename)

    try:
        run()
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()