nodes/1038lab-ComfyUI-RMBG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2025-11-25 13:53:09 -08:00
committed by GitHub
parent 148cffa236
commit df4dc75dda
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sam3/sam/transformer.py
+8 -5
View File
@@ -282,9 +282,8 @@ class RoPEAttention(Attention):
self.compute_cis = partial(
compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
)
device = torch.device("cuda") if torch.cuda.is_available() else None
self.freqs_cis = self.compute_cis(
end_x=feat_sizes[0], end_y=feat_sizes[1], device=device
end_x=feat_sizes[0], end_y=feat_sizes[1], device=None
)
if self.use_rope_real:
self.freqs_cis_real = self.freqs_cis.real
@@ -306,10 +305,14 @@ class RoPEAttention(Attention):
# Apply rotary position encoding
w = h = math.sqrt(q.shape[-2])
if self.freqs_cis.shape[0] != q.shape[-2]:
if (
self.freqs_cis.device != q.device
or self.freqs_cis.shape[0] != q.shape[-2]
):
self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device)
self.freqs_cis_real = self.freqs_cis.real
self.freqs_cis_imag = self.freqs_cis.imag
if self.use_rope_real:
self.freqs_cis_real = self.freqs_cis.real
self.freqs_cis_imag = self.freqs_cis.imag
if q.shape[-2] != k.shape[-2]:
assert self.rope_k_repeat