Init_Image like this...? / Suggestions for a better way?

[drscotthawley]

If we want to do init images / inpainting, is the following crude & ham-handed implementation sufficient (I mean, it works), or is there a more elegant solution?

Assuming we add extra args for "init_image" [-1...1] and an init_strength mask [0..1]:

(in sampling.py)

@torch.no_grad()
def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    sigmas_cpu = sigmas.detach().cpu().numpy()
    ds = []
    ### ADDED THE NEXT LINE
    modified_extra_args = {k: v for k, v in extra_args.items() if k not in ['init_mask','init_image']}

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **modified_extra_args)
        d = to_d(x, sigmas[i], denoised)
        ds.append(d)
        if len(ds) > order:
            ds.pop(0)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        cur_order = min(i + 1, order)
        coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
        x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
        #### ADDED THE NEXT FEW LINES
        t = i/len(sigmas)
        if 'init_image' in extra_args.keys() and extra_args['init_image'] is not None and \
            'init_mask' in extra_args.keys() and extra_args['init_mask'] is not None:
            x = torch.where(extra_args['init_mask'] >= t, extra_args['init_image'], x)

    return x

I found I had to do the modified_extra_args thing because the denoiser wasn't expecting those extra args.

The fact that this would mean editing every single sampling algorithm seems ...not good... but sufficient for my messing-around today.

Would love suggestions! Thanks.

[drscotthawley]

Ahhh wait, I should use the the callback... Ok.

        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})

Will close soon, once I figure out how to use this properly. ;-)

[drscotthawley]

Nope, I'm confused. My method worked in simple cases but for others the images just become all-black.
Although I'm told that maybe if I switch to ordinary DDPM sampling, inpainting may work better.

[drscotthawley]

Got it. :-) Closing. I'll put some deets in a paper soon. :-)