| 43 | |
| 44 | @register_to_config |
| 45 | def __init__( |
| 46 | self, |
| 47 | patch_size: int = 1, |
| 48 | in_channels: int = 64, |
| 49 | num_layers: int = 19, |
| 50 | num_single_layers: int = 38, |
| 51 | attention_head_dim: int = 128, |
| 52 | num_attention_heads: int = 24, |
| 53 | joint_attention_dim: int = 4096, |
| 54 | pooled_projection_dim: int = 768, |
| 55 | guidance_embeds: bool = False, |
| 56 | axes_dims_rope: List[int] = [16, 56, 56], |
| 57 | num_mode: int = None, |
| 58 | extra_conditioning_channels: int = 0, |
| 59 | extra_condition_channels: int = 0, |
| 60 | ): |
| 61 | super().__init__() |
| 62 | self.out_channels = in_channels |
| 63 | self.inner_dim = num_attention_heads * attention_head_dim |
| 64 | |
| 65 | self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope) |
| 66 | text_time_guidance_cls = ( |
| 67 | CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings |
| 68 | ) |
| 69 | self.time_text_embed = text_time_guidance_cls( |
| 70 | embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim |
| 71 | ) |
| 72 | |
| 73 | self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) |
| 74 | self.x_embedder = torch.nn.Linear(in_channels, self.inner_dim) |
| 75 | |
| 76 | self.transformer_blocks = nn.ModuleList( |
| 77 | [ |
| 78 | FluxTransformerBlock( |
| 79 | dim=self.inner_dim, |
| 80 | num_attention_heads=num_attention_heads, |
| 81 | attention_head_dim=attention_head_dim, |
| 82 | ) |
| 83 | for i in range(num_layers) |
| 84 | ] |
| 85 | ) |
| 86 | |
| 87 | self.single_transformer_blocks = nn.ModuleList( |
| 88 | [ |
| 89 | FluxSingleTransformerBlock( |
| 90 | dim=self.inner_dim, |
| 91 | num_attention_heads=num_attention_heads, |
| 92 | attention_head_dim=attention_head_dim, |
| 93 | ) |
| 94 | for i in range(num_single_layers) |
| 95 | ] |
| 96 | ) |
| 97 | |
| 98 | # controlnet_blocks |
| 99 | self.controlnet_blocks = nn.ModuleList([]) |
| 100 | for _ in range(len(self.transformer_blocks)): |
| 101 | self.controlnet_blocks.append(zero_module(nn.Linear(self.inner_dim, self.inner_dim))) |
| 102 | |