use_real set to False in get_frequencies

Author: prishajain1

src/maxdiffusion/models/embeddings_flax.py

@@ -244,10 +244,8 @@ def get_1d_rotary_pos_embed(
     freqs_cos = jnp.cos(freqs)
     freqs_sin = jnp.sin(freqs)
     out = jnp.stack([freqs_cos, -freqs_sin, freqs_sin, freqs_cos], axis=-1)
-    print("Using real rotary embeddings (Flux-style)")
   else:
     # Wan 2.1
-    print("Using complex rotary embeddings (Wan-style)")
     out = jnp.exp(1j * freqs)
   return out
 

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -46,7 +46,7 @@ def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int, use_r
   t_dim = attention_head_dim - h_dim - w_dim
   freqs = []
   for dim in [t_dim, h_dim, w_dim]:
-    freq = get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype=jnp.float32, use_real=use_real)
+    freq = get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype=jnp.float32, use_real=False)
     freqs.append(freq)
   freqs = jnp.concatenate(freqs, axis=1)
   t_size = attention_head_dim // 2 - 2 * (attention_head_dim // 6)