Spaces:
Running
on
Zero
Running
on
Zero
SunderAli17
commited on
Commit
•
5ac1546
1
Parent(s):
3fa6c11
Create math.py
Browse files- flux/math.py +31 -0
flux/math.py
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from einops import rearrange
|
3 |
+
from torch import Tensor
|
4 |
+
|
5 |
+
|
6 |
+
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
|
7 |
+
if pe is not None:
|
8 |
+
q, k = apply_rope(q, k, pe)
|
9 |
+
|
10 |
+
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
11 |
+
x = rearrange(x, "B H L D -> B L (H D)")
|
12 |
+
|
13 |
+
return x
|
14 |
+
|
15 |
+
|
16 |
+
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
|
17 |
+
assert dim % 2 == 0
|
18 |
+
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
|
19 |
+
omega = 1.0 / (theta**scale)
|
20 |
+
out = torch.einsum("...n,d->...nd", pos, omega)
|
21 |
+
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
|
22 |
+
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
|
23 |
+
return out.float()
|
24 |
+
|
25 |
+
|
26 |
+
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
|
27 |
+
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
28 |
+
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
29 |
+
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
|
30 |
+
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
|
31 |
+
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|