Which is the 'correct' way to add RePaint?

[drscotthawley]

Hi. I'm trying to enable RePaint in k-diffusion, but I'm not sure how exactly to convert from the original formulation of repaint, which uses (line 10 of Algorithm 1):

which I write line 10 in code as

x = torch.sqrt(1-beta)*x +  torch.sqrt(beta)*torch.randn_like(x)

But k-diffusion uses sigmas rather than betas, and these sigmas scale beyond 1.

Here are my various attempts to enable RePaint for one of the DPM++ samplers. Which of these -- if any, otherwise what else -- would be the 'correct' way to do this? Many of these variants produce NaNs when the number of repaint loops (here set via the repaint kwarg) exceeds more than a few.

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


@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 applies the mask (again?)
                    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 = 'var5' # ['orig','var1','var2', etc...]

                    sigma_diff = (sigmas[i] - sigmas[i+1]).abs()  # abs shouldn't be needed but...just in case
                    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 works "ok" for repaint=2, but larger values yield artifacts and/or NaNs.
                    elif repaint_choice == 'var1': # or maybe this...?  # ok too but worse than orig
                        x = x + sigma_diff*torch.randn_like(x)
                    elif repaint_choice == 'var2':  # or this...?    # yields NaNs pretty quick
                        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, very little 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