openai / consistency_models

Official repo for consistency models.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Code for Zero-Shot Image Editing

SDNAFIO opened this issue · comments

commented

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

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.

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?

Best wishes.

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

commented

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?

commented

@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

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.