Fix(backend): Denoise Masks being applied at the wrong time

[dunkeroni]

What type of PR is this? (check all applicable)

• Refactor
• Feature
• Bug Fix
• Optimization
• Documentation Update
• Community Node Submission

Have you discussed this change with the InvokeAI team?

• Yes
• No, because:

Have you updated all relevant documentation?

• Yes
• No

Description

Denoise Mask was being applied at the wrong time in the step loop, creating an off-by-one error. Each step during masked denoise was computing based on an unmasked region that was 1 step more noised than it should have been.

Related Tickets & Documents

• Closes [bug]: gradient denoising causes artifacts in combination with LCM #5966

QA Instructions, Screenshots, Recordings

This effect has been seen regularly in SDXL inpainting, since those models are better at matching the surrounding textures during denoising. It was most prevalent on smooth or otherwise featureless backgrounds. In anime models, it tends to present as sparkles or polkadots.

With very low step counts (8-step LCM) there was a fringe effect noticed on the edges of expanded masks.

Results after fix

Inpainting LCM still has some color drift, possibly affected by the VAE round trip. That can be investigated later on. The noise artifacting at least is gone.

Merge Plan

This PR can be merged when approved

Added/updated tests?

• Yes
• No : does not affect API for the diffusion pipeline

[blessedcoolant]

The results seem fine but the progress images not show that the entire image is being denoised rather than just the area. Is that intended?

[dunkeroni]

The results seem fine but the progress images not show that the entire image is being denoised rather than just the area. Is that intended?

It's expected, not intended. Previously the code was acting on all tensors in step_output with a special case to detect and replace the values in the "predicted original sample". The result of this process converges to the initial image, but the callback gets a preview that is updated each step.

We can add some extra code to capture this again if we want previews to be more stable.

Incidentally, it looks like some samplers call that output "denoised" so we weren't applying that logic to all samplers anyway.

[hipsterusername]

Think it’s a nice to have if we can get the preview back, but I think it should be a follow on. This is solving a big pain point on inpainting that I think worth getting in ASAP

[psychedelicious]

Incidentally, it looks like some samplers call that output "denoised" so we weren't applying that logic to all samplers anyway.

I've noticed this at some point and was wondering what was happening! Good to know.

@dunkeroni Barring a fix for the progress images, is there any reason to not merge this?

[dunkeroni]

@dunkeroni Barring a fix for the progress images, is there any reason to not merge this?

I was going to explain a framework for how I want to hijack the denoise_mask input on the node and generalize it into a denoise_guidance type that allows code injection into multiple points from other nodes... but for now here's a patch that makes the inpaint preview work again.

Also I was wrong. It's just LCM sampler that uses "denoised". A lot of the others have no field at all. This patch just takes whatever is there if there is something and otherwise uses the current latents to create a "pred_original_sample" that the progress callback can use.

To answer your question though, it is ready to merge.

[blessedcoolant]

Tested. Seems good. Works as advertised. Great job.

[psychedelicious]

Thanks dunkeroni