Self-Attention-Guidance (SAG) for ComfyUI

[ghost]

Have been searching for this. I always used FreeU together with SAG in A1111 to eliminate some of the artifacts that SD 1.5 and finetuned models bring. It can bring out some more texture and clean up strange artifacts in the background. Does it exist for ComfyUI?

[Slickytail]

Hello, I'm working on implementing this and would appreciate help from anyone. My fork is here.

I've never used SAG in A1111, so I'm not sure if the results I'm getting are right. Strictly speaking, everything works -- there's no crashes etc. However, it gets really ugly if I push the SAG guidance scale past 0.25. The things I'm not sure about:

• am I pulling the attention map from the right layer?
• am I correctly computing the estimate x0?

EDIT: everything works now.
If anyone wants to test this, the SAG node in my fork is _for_testing/Self-Attention Guidance.

[Slickytail]

@frankyifei SAG is 1 extra UNet call per diffusion step. Fooocus implements this by skipping the normal conditional step and simply incorporating the SAG blurring into the conditional step (see line 236 of modules/patch.py in fooocus codebase). This is a substantial modification to the recommended cfg+sag formula from the original paper. I have faith in Zhang if he thinks this works better, but in my comfy patch I chose to be more faithful to the paper.

[rollingcookies]

Hello, is there any way to modify the node to work in two modes - as it is now and as it is implemented in Fooocus?

[Slickytail]

@rollingcookies You would need to make some internal changes to the comfy sampling system. In comfy.samplers.sampling_function, you can skip the uncond pass, but not the cond pass. So you'll need to make a few changes there, as well as in comfy.samplers.calc_cond_uncond_batch in order to enable that.

(an easier way would just be to pass through the uncond as the cond, and then make SAG extract the attention map from the cond instead of the uncond, but this might break some other nodes/techniques that rely on modifying only the cond or only the uncond pass).

Then in the sag node, set whatever flag you created to skip the cond pass. Now in the post_cfg_function, just change the extra call to calc_cond_uncond_batch to pass through the cond instead. Finally, change the max scale in the SAG Node, because that becomes your new CFG scale.

[rollingcookies]

@Slickytail Thanks for your explanation, I would love to implement all this, but unfortunately I'm not a programmer, just a designer ;(

[rollingcookies]

@Slickytail I can help to test this node (and make examples of workflows) if you help write it!