Code for Zero-Shot Image Editing about consistency_models

[arittenbach]

Does anyone have an example of how to use the iterative inpainting function?

[aarontan-git]

@arittenbach I tried the following based on your recommendation but was not able to get the same as Figure 10 (mainly just got noise within the inpainted area.

Have you had any updates?

@aarontan-git Alright after playing around for awhile I think I figured it out. At the very least it generates results similar to what is seen in figure 10 in the paper using the LSUN bedroom LPIPS model

using the iterative inpaint function in cm/karras_diffusion.py :

def iterative_inpainting(
distiller, -> should be same model given to sample function in karras_sample function when generating imagers (variable denoiser)
images, -> stack of images to inpaint, tensor with (batch_size, 3, image_size, image_size), should be in range -1 to 1, i used the exact output of the karras_sample function
x, -> noise generated by generator, same size as image (batch_size, 3, image_size,image_size)
ts, -> from algorithm 4 in appendix of paper, ts should be range(2,41) if steps is 40
t_min=0.002, (i left this as is)
t_max=80.0, (i left this as is)
rho=7.0, (i left this as is)
steps=40, (i left this as is)
generator=None, (same generator used to generate images in karras_sample function)
):

otherwise, i used same parameters as in launch script for generating images using lsun_bedroom lpips with multistep.

hope this is clear

[aarontan-git]

Does anyone have an example of how to use the iterative inpainting function?

Also looking for the same - @arittenbach let me know if you've had any insight?

[ajrheng]

@jS5t3r In case you're still looking into this or for anyone else wondering, you should not pass distiller=model in iterative_inpainting. distiller should be the denoiser function in karras_sample, something like this:

        def denoiser(x_t, sigma):
            _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
            if args.clip_denoised:
                denoised = denoised.clamp(-1, 1)
            return denoised
    
        sample, original = iterative_inpainting(
            distiller=denoiser,
            images=ref_img,
            x=generator.randn(*ref_img.shape, device=dist_util.dev()),
            ts=ts,
            t_min=0.002,
            t_max=80.0,
            rho=diffusion.rho,
            steps=args.steps,
            generator=generator
        )

After doing this I get inpainting samples that resemble Fig 10 of the paper! Basically in your code, you pass in model directly, thus bypassing the skip connections (Eq. 5 of paper) that defines the Consistency Model that is defined in diffusion.denoise, which is incorrect and leads to artifacts.

[sharkDDD]

Is it planned to release the code for the Zero-Shot Image Editing? In particular, I'd be interested in the super-resolution.

Best regards

The function iterative_colorization/inpainting/superres are contained in karras_diffusion.py, I think maybe this is what you are looking for.

[pzs19]

thanks for your code in karras_diffusion.py. But I've noticed that the effect of the superres heavily rely on some hyperparameters when using cd_bedroom256_l2.pt, say ts in https://github.com/openai/consistency_models/blob/6d26080c58244555c031dbc63080c0961af74200/cm/karras_diffusion.py#LL838C8-L838C8
So are there some recommended hyperparameter for superres? or a scripts like launsh.sh?

[pzs19]

Best wishes.

[jS5t3r]

Launch File:

mpiexec -n 1 python image_sample_inpainting.py --batch_size 32 \
--training_mode consistency_distillation \
--sampler multistep \
--ts 2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40 \
--steps 40 \
--model_path checkpoints/cd_bedroom256_lpips.pt \
--attention_resolutions 32,16,8 --class_cond False --use_scale_shift_norm False --dropout 0.0 \
--image_size 256 --num_channels 256 --num_head_channels 64 --num_res_blocks 2 \
--num_samples 32 \
--resblock_updown True --use_fp16 True --weight_schedule uniform

My code looks like this:

    model_kwargs = {}

    sample = karras_sample(
        diffusion,
        model,
        (args.batch_size, 3, args.image_size, args.image_size),
        steps=args.steps,
        model_kwargs=model_kwargs,
        device=dist_util.dev(),
        clip_denoised=args.clip_denoised,
        sampler=args.sampler,
        sigma_min=args.sigma_min,
        sigma_max=args.sigma_max,
        s_churn=args.s_churn,
        s_tmin=args.s_tmin,
        s_tmax=args.s_tmax,
        s_noise=args.s_noise,
        generator=generator,
        ts=ts,
    )

    x_out, sample = iterative_inpainting(
        distiller=model,
        images=sample,
        x=generator.randn(*sample.shape, device=dist_util.dev()), #* sigma_max
        ts=ts,
        t_min=0.002,
        t_max=80.0,
        rho=7.0,
        steps=40,
        generator=generator,
    )

    x_out = ((x_out + 1) * 127.5).clamp(0, 255).to(th.uint8)
    x_out = x_out.permute(0, 2, 3, 1)
    x_out = x_out.contiguous()

    sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
    sample = sample.permute(0, 2, 3, 1)
    sample = sample.contiguous()

The masked images look ok.

@aarontan-git @arittenbach Could you find some improvement? The inpainting does not work properly, what am I doing wrong?

In Algorithm 4 in the paper, they described A as an invertible linear transformation, that maps images to the latent space.
I cannot see any mapping to the latent space.