Implement NAG on all the models based on the Flux code. (#12500) · code/app-python-comfyui@1892753 · GitHub
Skip to content

Commit 1892753

Browse files
comfyanonymouscomfyanonymous
authored
Implement NAG on all the models based on the Flux code. (Comfy-Org#12500)
Use the Normalized Attention Guidance node. Flux, Flux2, Klein, Chroma, Chroma radiance, Hunyuan Video, etc..
1 parent 8a6fbc2 commit 1892753

7 files changed

Lines changed: 128 additions & 1 deletion

File tree

‎comfy/ldm/chroma/model.py‎

Lines changed: 2 additions & 0 deletions

‎comfy/ldm/flux/layers.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,9 @@ def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=N
196196
else:
197197
(img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec
198198

199+
transformer_patches = transformer_options.get("patches", {})
200+
extra_options = transformer_options.copy()
201+
199202
# prepare image for attention
200203
img_modulated = self.img_norm1(img)
201204
img_modulated = apply_mod(img_modulated, (1 + img_mod1.scale), img_mod1.shift, modulation_dims_img)
@@ -224,6 +227,12 @@ def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=N
224227
attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
225228
del q, k, v
226229

230+
if "attn1_output_patch" in transformer_patches:
231+
extra_options["img_slice"] = [txt.shape[1], attn.shape[1]]
232+
patch = transformer_patches["attn1_output_patch"]
233+
for p in patch:
234+
attn = p(attn, extra_options)
235+
227236
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]
228237

229238
# calculate the img bloks
@@ -303,6 +312,9 @@ def forward(self, x: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation
303312
else:
304313
mod = vec
305314

315+
transformer_patches = transformer_options.get("patches", {})
316+
extra_options = transformer_options.copy()
317+
306318
qkv, mlp = torch.split(self.linear1(apply_mod(self.pre_norm(x), (1 + mod.scale), mod.shift, modulation_dims)), [3 * self.hidden_size, self.mlp_hidden_dim_first], dim=-1)
307319

308320
q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
@@ -312,6 +324,12 @@ def forward(self, x: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation
312324
# compute attention
313325
attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
314326
del q, k, v
327+
328+
if "attn1_output_patch" in transformer_patches:
329+
patch = transformer_patches["attn1_output_patch"]
330+
for p in patch:
331+
attn = p(attn, extra_options)
332+
315333
# compute activation in mlp stream, cat again and run second linear layer
316334
if self.yak_mlp:
317335
mlp = self.mlp_act(mlp[..., self.mlp_hidden_dim_first // 2:]) * mlp[..., :self.mlp_hidden_dim_first // 2]

‎comfy/ldm/flux/model.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ def forward_orig(
142142
attn_mask: Tensor = None,
143143
) -> Tensor:
144144

145+
transformer_options = transformer_options.copy()
145146
patches = transformer_options.get("patches", {})
146147
patches_replace = transformer_options.get("patches_replace", {})
147148
if img.ndim != 3 or txt.ndim != 3:
@@ -231,6 +232,7 @@ def block_wrap(args):
231232

232233
transformer_options["total_blocks"] = len(self.single_blocks)
233234
transformer_options["block_type"] = "single"
235+
transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
234236
for i, block in enumerate(self.single_blocks):
235237
transformer_options["block_index"] = i
236238
if ("single_block", i) in blocks_replace:

‎comfy/ldm/hunyuan_video/model.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,7 @@ def forward_orig(
304304
control=None,
305305
transformer_options={},
306306
) -> Tensor:
307+
transformer_options = transformer_options.copy()
307308
patches_replace = transformer_options.get("patches_replace", {})
308309

309310
initial_shape = list(img.shape)
@@ -416,6 +417,7 @@ def block_wrap(args):
416417

417418
transformer_options["total_blocks"] = len(self.single_blocks)
418419
transformer_options["block_type"] = "single"
420+
transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
419421
for i, block in enumerate(self.single_blocks):
420422
transformer_options["block_index"] = i
421423
if ("single_block", i) in blocks_replace:

‎comfy/model_patcher.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -406,13 +406,16 @@ def clone_has_same_weights(self, clone: 'ModelPatcher'):
406406
def memory_required(self, input_shape):
407407
return self.model.memory_required(input_shape=input_shape)
408408

409+
def disable_model_cfg1_optimization(self):
410+
self.model_options["disable_cfg1_optimization"] = True
411+
409412
def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
410413
if len(inspect.signature(sampler_cfg_function).parameters) == 3:
411414
self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
412415
else:
413416
self.model_options["sampler_cfg_function"] = sampler_cfg_function
414417
if disable_cfg1_optimization:
415-
self.model_options["disable_cfg1_optimization"] = True
418+
self.disable_model_cfg1_optimization()
416419

417420
def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
418421
self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)

‎comfy_extras/nodes_nag.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import torch
2+
from comfy_api.latest import ComfyExtension, io
3+
from typing_extensions import override
4+
5+
6+
class NAGuidance(io.ComfyNode):
7+
@classmethod
8+
def define_schema(cls) -> io.Schema:
9+
return io.Schema(
10+
node_id="NAGuidance",
11+
display_name="Normalized Attention Guidance",
12+
description="Applies Normalized Attention Guidance to models, enabling negative prompts on distilled/schnell models.",
13+
category="",
14+
is_experimental=True,
15+
inputs=[
16+
io.Model.Input("model", tooltip="The model to apply NAG to."),
17+
io.Float.Input("nag_scale", min=0.0, default=5.0, max=50.0, step=0.1, tooltip="The guidance scale factor. Higher values push further from the negative prompt."),
18+
io.Float.Input("nag_alpha", min=0.0, default=0.5, max=1.0, step=0.01, tooltip="Blending factor for the normalized attention. 1.0 is full replacement, 0.0 is no effect."),
19+
io.Float.Input("nag_tau", min=1.0, default=1.5, max=10.0, step=0.01),
20+
# io.Float.Input("start_percent", min=0.0, default=0.0, max=1.0, step=0.01, tooltip="The relative sampling step to begin applying NAG."),
21+
# io.Float.Input("end_percent", min=0.0, default=1.0, max=1.0, step=0.01, tooltip="The relative sampling step to stop applying NAG."),
22+
],
23+
outputs=[
24+
io.Model.Output(tooltip="The patched model with NAG enabled."),
25+
],
26+
)
27+
28+
@classmethod
29+
def execute(cls, model: io.Model.Type, nag_scale: float, nag_alpha: float, nag_tau: float) -> io.NodeOutput:
30+
m = model.clone()
31+
32+
# sigma_start = m.get_model_object("model_sampling").percent_to_sigma(start_percent)
33+
# sigma_end = m.get_model_object("model_sampling").percent_to_sigma(end_percent)
34+
35+
def nag_attention_output_patch(out, extra_options):
36+
cond_or_uncond = extra_options.get("cond_or_uncond", None)
37+
if cond_or_uncond is None:
38+
return out
39+
40+
if not (1 in cond_or_uncond and 0 in cond_or_uncond):
41+
return out
42+
43+
# sigma = extra_options.get("sigmas", None)
44+
# if sigma is not None and len(sigma) > 0:
45+
# sigma = sigma[0].item()
46+
# if sigma > sigma_start or sigma < sigma_end:
47+
# return out
48+
49+
img_slice = extra_options.get("img_slice", None)
50+
51+
if img_slice is not None:
52+
orig_out = out
53+
out = out[:, img_slice[0]:img_slice[1]] # only apply on img part
54+
55+
batch_size = out.shape[0]
56+
half_size = batch_size // len(cond_or_uncond)
57+
58+
ind_neg = cond_or_uncond.index(1)
59+
ind_pos = cond_or_uncond.index(0)
60+
z_pos = out[half_size * ind_pos:half_size * (ind_pos + 1)]
61+
z_neg = out[half_size * ind_neg:half_size * (ind_neg + 1)]
62+
63+
guided = z_pos * nag_scale - z_neg * (nag_scale - 1.0)
64+
65+
eps = 1e-6
66+
norm_pos = torch.norm(z_pos, p=1, dim=-1, keepdim=True).clamp_min(eps)
67+
norm_guided = torch.norm(guided, p=1, dim=-1, keepdim=True).clamp_min(eps)
68+
69+
ratio = norm_guided / norm_pos
70+
scale_factor = torch.minimum(ratio, torch.full_like(ratio, nag_tau)) / ratio
71+
72+
guided_normalized = guided * scale_factor
73+
74+
z_final = guided_normalized * nag_alpha + z_pos * (1.0 - nag_alpha)
75+
76+
if img_slice is not None:
77+
orig_out[half_size * ind_neg:half_size * (ind_neg + 1), img_slice[0]:img_slice[1]] = z_final
78+
orig_out[half_size * ind_pos:half_size * (ind_pos + 1), img_slice[0]:img_slice[1]] = z_final
79+
return orig_out
80+
else:
81+
out[half_size * ind_pos:half_size * (ind_pos + 1)] = z_final
82+
return out
83+
84+
m.set_model_attn1_output_patch(nag_attention_output_patch)
85+
m.disable_model_cfg1_optimization()
86+
87+
return io.NodeOutput(m)
88+
89+
90+
class NagExtension(ComfyExtension):
91+
@override
92+
async def get_node_list(self) -> list[type[io.ComfyNode]]:
93+
return [
94+
NAGuidance,
95+
]
96+
97+
98+
async def comfy_entrypoint() -> NagExtension:
99+
return NagExtension()

‎nodes.py‎

Lines changed: 1 addition & 0 deletions

0 commit comments

Comments
 (0)