Adding 'strength' parameter to StableDiffusionInpaintingPipeline

[rupertmenneer]

Currently in inpainting pipelines (e.g. StableDiffusionInpaintingPipeline), masked areas are expected to be 100% noise. During the diffussion sampling process the masked areas are then generated w.r.t the visible image. The problem with this process is that it assumes you have no estimate of the underlying masked area such as a rough sketch. This means you have little control other than the text prompt to control the output, in particular structure and colour are difficult to explicitly control.

I propose a parameter to be passed into inpainting pipelines that allow you to skip to a specific point in the sampling process. The parameter will be called skip_time_percentage e.g. if 0.75 is passed then 75% of the sampling steps are skipped. The image and masked regions are initialised by applying forward diffusion to that point in diffusion time. This means the input is actually a noisy version of the initial image rather than pure noise in masked regions.

I have used this before to great success and of course it also speeds up sampling time by skipping a large portion of steps.

This feature has been integrated into the Imagen diffusion pytorch implementation here for some time: https://github.com/lucidrains/imagen-pytorch

Is this something that people would be interested in? I can write up the code and submit a PR if so.

[rupertmenneer]

I just had a look at the above feature, to confirm the sort of technique I am describing see below

Example input (original, noisy masked with 80% skip, mask):

It wasn't totally clear to me how to initialise the stable diffusion inpainting pipeline with the above input. It seems there are currently the following inputs into the model.

• latents - initialised random gaussian noise which is scaled by the samplers init_noise_sigma, this is set to 1 in the PNDMScheduler. This random noise has the same shape as the image after VAE encoding.
• mask - a binary mask that has been passed through the VAE encoder.
• masked_image_latents - the original image multiplied by the mask (e.g. 0 in regions set to be inpainted), this is also been passed through the VAE encoder.

I suppose the question to answer is how to correctly initialise the pipeline halfway through the process. I tried the following:

# Handling less timesteps
            # 5. set timesteps
            self.scheduler.set_timesteps(num_inference_steps, device=device)
            timesteps = self.scheduler.timesteps
            # num_inference_steps
            adjusted_num_steps = int(num_inference_steps * skip_time_percentage)
            timesteps = timesteps[adjusted_num_steps:]

# Latent initialisation
            # instead of initialising with pure noise, create a noisy version of the input and convert to latent space
            noisy_input = masked_image.clone().to(device=device, dtype=dtype)
            noise = randn_tensor(img_shape, generator=generator, device=device, dtype=dtype)
            noisy_input = self.scheduler.add_noise(noisy_input, noise, torch.tensor(max_timestep, dtype=dtype, device=device).long())
            
            new_latents = self.vae.encode(noisy_input).latent_dist.sample(generator=generator)
            new_latents = self.vae.config.scaling_factor * new_latents


        

Essentially initialising the latents with a noisy version of the input instead of with pure gaussian noise, and then starting the sampling process from a 'midpoint' rather than from one end. A few other adjustments such as how warmup steps are calculated/when the image is masked etc to make the above changes logical.

However this was not successful and seems to have a fairly different distribution to pure noise even when the max timestep is set to the number of timesteps in the sampler! Any insight on this would be useful as I haven't played too much with latent space diffusion or samplers that produce forward diffusion outputs that aren't very similar to gaussian noise at the maximum timestep. Hope that makes sense! Will put this on ice for now

CC: @patrickvonplaten @williamberman

[williamberman]

hey @rupertmenneer thanks so much for the issue!

I think we do have this parameter in other pipelines, called strength

diffusers/src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py

Lines 874 to 879 in 5c7a35a

	strength (`float`, *optional*, defaults to 0.8):
	Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
	will be used as a starting point, adding more noise to it the larger the `strength`. The number of
	denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
	be maximum and the denoising process will run for the full number of iterations specified in
	`num_inference_steps`. A value of 1, therefore, essentially ignores `image`.

I agree it would be worth adding it to the pipelines you mentioned that don't support it rn. Feel free to open a PR and tag me, happy to review :)

[rupertmenneer]

hey @williamberman that is super helpful! and exactly what I was looking for. Interestingly, it looks exactly like what I was trying but without success on the stable diffusion inpainting pipeline, I will review my code and see if I can find a culprit.

So yes, I will aim to write a PR adding a strength to the StableDiffusionInpaintingPipeline and other relevant ones when I get a chance :) thanks again!

[rupertmenneer]

Found some time today to quickly add the strength parameter into the StableDiffusionInpaintingPipeline

Got some pretty cool results! Going to add a unit test and create a PR for this now.

strengths = [1., 0.75, 0.5, 0.25, 0.]
prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
for strength in reversed(strengths):
    output_image = pipe(prompt=prompt,
                image=init_image,
                mask_image=mask_image,
                strength=strength,
                num_inference_steps=50).images[0]

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.