Parameterize frequency_embedding_size in CombinedTimestepGuidanceTextProjEmbeddings

Author: amepas

src/maxdiffusion/models/embeddings_flax.py

@@ -497,20 +497,21 @@ class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
   embedding_dim: int
   pooled_projection_dim: int
   guidance_embeds: bool = True
+  frequency_embedding_size: int = 256
   dtype: jnp.dtype = jnp.float32
   weights_dtype: jnp.dtype = jnp.float32
   precision: jax.lax.Precision = None
 
   @nn.compact
   def __call__(self, timestep, guidance, pooled_projection=None):
-    timesteps_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(timestep)
+    timesteps_proj = FlaxTimesteps(dim=self.frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0)(timestep)
     dtype = pooled_projection.dtype if pooled_projection is not None else jnp.float32
     timestep_emb = FlaxTimestepEmbedding(
         time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype
     )(timesteps_proj.astype(dtype))
 
     if self.guidance_embeds and guidance is not None:
-      guidance_proj = FlaxTimesteps(dim=256, flip_sin_to_cos=True, freq_shift=0)(guidance)
+      guidance_proj = FlaxTimesteps(dim=self.frequency_embedding_size, flip_sin_to_cos=True, freq_shift=0)(guidance)
       guidance_emb = FlaxTimestepEmbedding(
           time_embed_dim=self.embedding_dim, dtype=self.dtype, weights_dtype=self.weights_dtype
       )(guidance_proj.astype(dtype))