ComfyUI-NAG/utils.py at main · Delcado19/ComfyUI-NAG · GitHub

1
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import math
import torch


def nag(z_positive, z_negative, scale, tau, alpha):
    m = min(z_positive.shape[0], z_negative.shape[0])
    if m == 0:
        return z_negative
    z_positive = z_positive[-m:]
    z_negative = z_negative[-m:]
    z_guidance = z_positive * scale - z_negative * (scale - 1)

    eps = 1e-6
    norm_positive = (
        torch.norm(z_positive, p=1, dim=-1, keepdim=True)
        .clamp_min(eps)
        .expand_as(z_positive)
    )
    norm_guidance = (
        torch.norm(z_guidance, p=1, dim=-1, keepdim=True)
        .clamp_min(eps)
        .expand_as(z_guidance)
    )

    s = norm_guidance / norm_positive
    z_guidance = z_guidance * torch.minimum(s, s.new_full((1,), tau)) / s

    z_guidance = z_guidance * alpha + z_positive * (1 - alpha)

    return z_guidance


def cat_context(context, nag_negative_context, trim_context=False, dim=1):
    assert dim in [1, 2]
    nag_negative_context = nag_negative_context.to(context)

    context_len = context.shape[dim]
    nag_neg_context_len = nag_negative_context.shape[dim]

    if context_len < nag_neg_context_len:
        if dim == 1:
            context = context.repeat(1, math.ceil(nag_neg_context_len / context_len), 1)
            if trim_context:
                context = context[:, -nag_neg_context_len:]
        else:
            context = context.repeat(1, 1, math.ceil(nag_neg_context_len / context_len), 1)
            if trim_context:
                context = context[:, :, -nag_neg_context_len:]

        context_len = context.shape[dim]

    if dim == 1:
        nag_negative_context = nag_negative_context.repeat(1, math.ceil(context_len / nag_neg_context_len), 1)
        nag_negative_context = nag_negative_context[:, -context_len:]
    else:
        nag_negative_context = nag_negative_context.repeat(1, 1, math.ceil(context_len / nag_neg_context_len), 1)
        nag_negative_context = nag_negative_context[:, :, -context_len:]


    return torch.cat([context, nag_negative_context], dim=0)


def check_nag_activation(transformer_options, nag_sigma_end):
    apply_nag = torch.all(transformer_options["sigmas"] >= nag_sigma_end)
    positive_batch = 0 in transformer_options["cond_or_uncond"]
    return apply_nag and positive_batch


def get_closure_vars(func):
    if func.__closure__ is None:
        return {}
    return {
        var: cell.cell_contents
        for var, cell in zip(func.__code__.co_freevars, func.__closure__)
    }


def is_from_wavespeed(func):
    closure = get_closure_vars(func)
    return "residual_diff_threshold" in closure \
        and "validate_can_use_cache_function" in closure


class NAGSwitch:
    def __init__(
        self,
        model: torch.nn.Module,
        nag_negative_cond,
        nag_scale, nag_tau, nag_alpha, nag_sigma_end,
    ):
        self.model = model
        self.nag_negative_cond = nag_negative_cond
        self.nag_scale = nag_scale
        self.nag_tau = nag_tau
        self.nag_alpha = nag_alpha
        self.nag_sigma_end = nag_sigma_end
        self.origin_forward = model.forward

    def set_nag(self):
        pass

    def set_origin(self):
        self.model.forward = self.origin_forward


# https://github.com/welltop-cn/ComfyUI-TeaCache/blob/4bca908bf53b029ea5739cb69ef2a9e6c06e6752/nodes.py
def poly1d(coefficients, x):
    result = torch.zeros_like(x)
    for i, coeff in enumerate(coefficients):
        result += coeff * (x ** (len(coefficients) - 1 - i))
    return result