Update transformer_bria.py

transformer_bria.py

@@ -10,38 +10,18 @@ from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_l
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.embeddings import TimestepEmbedding, get_timestep_embedding
 from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
-
-# Support different diffusers versions
-try:
-    from diffusers.models.embeddings import FluxPosEmbed as EmbedND
-except:
-    from diffusers.models.transformers.transformer_flux import rope
-    class EmbedND(nn.Module):
-        def __init__(self, theta: int, axes_dim: List[int]):
-            super().__init__()
-            self.theta = theta
-            self.axes_dim = axes_dim
-
-        def forward(self, ids: torch.Tensor) -> torch.Tensor:
-            n_axes = ids.shape[-1]
-            emb = torch.cat(
-                [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
-                dim=-3,
-            )
-            return emb.unsqueeze(1)
-
-
+from bria_utils import FluxPosEmbed as EmbedND
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 class Timesteps(nn.Module):
-    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1,max_period=10000):
+    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1,time_theta=10000):
         super().__init__()
         self.num_channels = num_channels
         self.flip_sin_to_cos = flip_sin_to_cos
         self.downscale_freq_shift = downscale_freq_shift
         self.scale = scale
-        self.max_period=max_period
+        self.time_theta=time_theta
 
     def forward(self, timesteps):
         t_emb = get_timestep_embedding(
@@ -50,15 +30,15 @@ class Timesteps(nn.Module):
             flip_sin_to_cos=self.flip_sin_to_cos,
             downscale_freq_shift=self.downscale_freq_shift,
             scale=self.scale,
-            max_period=self.max_period
+            max_period=self.time_theta
         )
         return t_emb
     
 class TimestepProjEmbeddings(nn.Module):
-    def __init__(self, embedding_dim, max_period):
+    def __init__(self, embedding_dim, time_theta):
         super().__init__()
 
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0,max_period=max_period)
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0,time_theta=time_theta)
         self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
         
     def forward(self, timestep, dtype):
@@ -106,7 +86,7 @@ class BriaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrig
         guidance_embeds: bool = False,
         axes_dims_rope: List[int] = [16, 56, 56],
         rope_theta = 10000,
-        max_period = 10000
+        time_theta = 10000
     ):
         super().__init__()
         self.out_channels = in_channels
@@ -116,7 +96,7 @@ class BriaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrig
         
         
         self.time_embed = TimestepProjEmbeddings(
-            embedding_dim=self.inner_dim,max_period=max_period
+            embedding_dim=self.inner_dim,time_theta=time_theta
         )
 
         # if pooled_projection_dim: