INNER CODE UNIT · Python
attention_basic_with_temperature
Extraltodeus/ComfyUI-AutomaticCFG · experimental_temperature.py:22
def attention_basic_with_temperature(self, q, k, v, extra_options, mask=None, attn_precision=None):
if isinstance(extra_options, int):
heads = extra_options
else:
heads = extra_options['n_heads']
b, _, dim_head = q.shape
dim_head //= heads
scale = dim_head ** -0.5
h = heads
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, -1, heads, dim_head)
.permute(0, 2, 1, 3)
.reshape(b * heads, -1, dim_head)
.contiguous(),
(q, k, v),