openai / consistency_models

Official repo for consistency models.
MIT License
6.08k stars 411 forks source link

Code for Zero-Shot Image Editing #5

Open SDNAFIO opened 1 year ago

SDNAFIO commented 1 year ago

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

sharkDDD commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

Best wishes.

arittenbach commented 1 year ago

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

aarontan-git commented 1 year ago

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?

aarontan-git commented 1 year ago

@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
jS5t3r commented 11 months ago

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.

ajrheng commented 4 months ago

@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.