CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

llm

Browse Source
This commit is contained in:
Haoming 2026-02-13 13:44:05 +08:00
parent ef0de020e3
commit dfa4474cc1
1 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
backend/nn/llm/llama.py
View File

@ -1,4 +1,4 @@
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.75/comfy/text_encoders/llama.py
# https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/text_encoders/llama.py
import math
from dataclasses import asdict, dataclass
@ -133,13 +133,6 @@ class Gemma2_2B_Config:
final_norm: bool = True
def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_dims=None, device=None):
if not isinstance(theta, list):
theta = [theta]
@ -168,7 +161,9 @@ def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_di
else:
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
out.append((cos, sin))
sin_split = sin.shape[-1] // 2
out.append((cos, sin[..., :sin_split], -sin[..., sin_split:]))
if len(out) == 1:
return out[0]
@ -180,8 +175,18 @@ def apply_rope(xq, xk, freqs_cis):
org_dtype = xq.dtype
cos = freqs_cis[0]
sin = freqs_cis[1]
q_embed = (xq * cos) + (rotate_half(xq) * sin)
k_embed = (xk * cos) + (rotate_half(xk) * sin)
nsin = freqs_cis[2]
q_embed = xq * cos
q_split = q_embed.shape[-1] // 2
q_embed[..., :q_split].addcmul_(xq[..., q_split:], nsin)
q_embed[..., q_split:].addcmul_(xq[..., :q_split], sin)
k_embed = xk * cos
k_split = k_embed.shape[-1] // 2
k_embed[..., :k_split].addcmul_(xk[..., k_split:], nsin)
k_embed[..., k_split:].addcmul_(xk[..., :k_split], sin)
return q_embed.to(org_dtype), k_embed.to(org_dtype)