(
self,
in_channels: int = 32,
out_channels: Optional[int] = 32,
num_attention_heads: int = 70,
attention_head_dim: int = 32,
num_layers: int = 20,
num_cross_attention_heads: Optional[int] = 20,
cross_attention_head_dim: Optional[int] = 112,
cross_attention_dim: Optional[int] = 2240,
caption_channels: int = 2304,
mlp_ratio: float = 2.5,
dropout: float = 0.0,
attention_bias: bool = False,
sample_size: int = 32,
patch_size: int = 1,
norm_elementwise_affine: bool = False,
norm_eps: float = 1e-6,
interpolation_scale: Optional[int] = None,
guidance_embeds: bool = False,
guidance_embeds_scale: float = 0.1,
qk_norm: Optional[str] = None,
timestep_scale: float = 1.0,
)
| 342 | |
| 343 | @register_to_config |
| 344 | def __init__( |
| 345 | self, |
| 346 | in_channels: int = 32, |
| 347 | out_channels: Optional[int] = 32, |
| 348 | num_attention_heads: int = 70, |
| 349 | attention_head_dim: int = 32, |
| 350 | num_layers: int = 20, |
| 351 | num_cross_attention_heads: Optional[int] = 20, |
| 352 | cross_attention_head_dim: Optional[int] = 112, |
| 353 | cross_attention_dim: Optional[int] = 2240, |
| 354 | caption_channels: int = 2304, |
| 355 | mlp_ratio: float = 2.5, |
| 356 | dropout: float = 0.0, |
| 357 | attention_bias: bool = False, |
| 358 | sample_size: int = 32, |
| 359 | patch_size: int = 1, |
| 360 | norm_elementwise_affine: bool = False, |
| 361 | norm_eps: float = 1e-6, |
| 362 | interpolation_scale: Optional[int] = None, |
| 363 | guidance_embeds: bool = False, |
| 364 | guidance_embeds_scale: float = 0.1, |
| 365 | qk_norm: Optional[str] = None, |
| 366 | timestep_scale: float = 1.0, |
| 367 | ) -> None: |
| 368 | super().__init__() |
| 369 | |
| 370 | out_channels = out_channels or in_channels |
| 371 | inner_dim = num_attention_heads * attention_head_dim |
| 372 | |
| 373 | # 1. Patch Embedding |
| 374 | self.patch_embed = PatchEmbed( |
| 375 | height=sample_size, |
| 376 | width=sample_size, |
| 377 | patch_size=patch_size, |
| 378 | in_channels=in_channels, |
| 379 | embed_dim=inner_dim, |
| 380 | interpolation_scale=interpolation_scale, |
| 381 | pos_embed_type="sincos" if interpolation_scale is not None else None, |
| 382 | ) |
| 383 | |
| 384 | # 2. Additional condition embeddings |
| 385 | if guidance_embeds: |
| 386 | self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) |
| 387 | self.time_embed_t = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) |
| 388 | self.time_embed_r = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) |
| 389 | |
| 390 | else: |
| 391 | self.time_embed = AdaLayerNormSingle(inner_dim) |
| 392 | self.time_embed_t = AdaLayerNormSingle(inner_dim) |
| 393 | self.time_embed_r = AdaLayerNormSingle(inner_dim) |
| 394 | |
| 395 | |
| 396 | self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) |
| 397 | self.caption_norm = RMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) |
| 398 | |
| 399 | # 3. Transformer blocks |
| 400 | self.transformer_blocks = nn.ModuleList( |
| 401 | [ |
nothing calls this directly
no test coverage detected