Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
renzhc
diffusers_dcu
Commits
3a31b291
Unverified
Commit
3a31b291
authored
Jun 02, 2025
by
Roy Hvaara
Committed by
GitHub
Jun 02, 2025
Browse files
Use float32 RoPE freqs in Wan with MPS backends (#11643)
Use float32 for RoPE on MPS in Wan
parent
b975bcef
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
4 additions
and
2 deletions
+4
-2
src/diffusers/models/transformers/transformer_wan.py
src/diffusers/models/transformers/transformer_wan.py
+4
-2
No files found.
src/diffusers/models/transformers/transformer_wan.py
View file @
3a31b291
...
@@ -72,7 +72,8 @@ class WanAttnProcessor2_0:
...
@@ -72,7 +72,8 @@ class WanAttnProcessor2_0:
if
rotary_emb
is
not
None
:
if
rotary_emb
is
not
None
:
def
apply_rotary_emb
(
hidden_states
:
torch
.
Tensor
,
freqs
:
torch
.
Tensor
):
def
apply_rotary_emb
(
hidden_states
:
torch
.
Tensor
,
freqs
:
torch
.
Tensor
):
x_rotated
=
torch
.
view_as_complex
(
hidden_states
.
to
(
torch
.
float64
).
unflatten
(
3
,
(
-
1
,
2
)))
dtype
=
torch
.
float32
if
hidden_states
.
device
.
type
==
"mps"
else
torch
.
float64
x_rotated
=
torch
.
view_as_complex
(
hidden_states
.
to
(
dtype
).
unflatten
(
3
,
(
-
1
,
2
)))
x_out
=
torch
.
view_as_real
(
x_rotated
*
freqs
).
flatten
(
3
,
4
)
x_out
=
torch
.
view_as_real
(
x_rotated
*
freqs
).
flatten
(
3
,
4
)
return
x_out
.
type_as
(
hidden_states
)
return
x_out
.
type_as
(
hidden_states
)
...
@@ -190,9 +191,10 @@ class WanRotaryPosEmbed(nn.Module):
...
@@ -190,9 +191,10 @@ class WanRotaryPosEmbed(nn.Module):
t_dim
=
attention_head_dim
-
h_dim
-
w_dim
t_dim
=
attention_head_dim
-
h_dim
-
w_dim
freqs
=
[]
freqs
=
[]
freqs_dtype
=
torch
.
float32
if
torch
.
backends
.
mps
.
is_available
()
else
torch
.
float64
for
dim
in
[
t_dim
,
h_dim
,
w_dim
]:
for
dim
in
[
t_dim
,
h_dim
,
w_dim
]:
freq
=
get_1d_rotary_pos_embed
(
freq
=
get_1d_rotary_pos_embed
(
dim
,
max_seq_len
,
theta
,
use_real
=
False
,
repeat_interleave_real
=
False
,
freqs_dtype
=
torch
.
float64
dim
,
max_seq_len
,
theta
,
use_real
=
False
,
repeat_interleave_real
=
False
,
freqs_dtype
=
freqs_dtype
)
)
freqs
.
append
(
freq
)
freqs
.
append
(
freq
)
self
.
freqs
=
torch
.
cat
(
freqs
,
dim
=
1
)
self
.
freqs
=
torch
.
cat
(
freqs
,
dim
=
1
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment