Unverified Commit 3335e226 authored by SahilCarterr's avatar SahilCarterr Committed by GitHub
Browse files

[FIX] Bug in FluxPosEmbed (#10115)



* Fix get_1d_rotary_pos_embed in embedding.py

* Update embeddings.py

---------
Co-authored-by: default avatarhlky <hlky@hlky.ac>
parent 65ab1052
......@@ -959,7 +959,12 @@ class FluxPosEmbed(nn.Module):
freqs_dtype = torch.float32 if is_mps else torch.float64
for i in range(n_axes):
cos, sin = get_1d_rotary_pos_embed(
self.axes_dim[i], pos[:, i], repeat_interleave_real=True, use_real=True, freqs_dtype=freqs_dtype
self.axes_dim[i],
pos[:, i],
theta=self.theta,
repeat_interleave_real=True,
use_real=True,
freqs_dtype=freqs_dtype,
)
cos_out.append(cos)
sin_out.append(sin)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment