(self, disable_guidance_embedder=False, input_dim=64, num_blocks=19)
| 277 | |
| 278 | class FluxDiT(torch.nn.Module): |
| 279 | def __init__(self, disable_guidance_embedder=False, input_dim=64, num_blocks=19): |
| 280 | super().__init__() |
| 281 | self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56]) |
| 282 | self.time_embedder = TimestepEmbeddings(256, 3072) |
| 283 | self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072) |
| 284 | self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072)) |
| 285 | self.context_embedder = torch.nn.Linear(4096, 3072) |
| 286 | self.x_embedder = torch.nn.Linear(input_dim, 3072) |
| 287 | |
| 288 | self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_blocks)]) |
| 289 | self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(38)]) |
| 290 | |
| 291 | self.final_norm_out = AdaLayerNormContinuous(3072) |
| 292 | self.final_proj_out = torch.nn.Linear(3072, 64) |
| 293 | |
| 294 | self.input_dim = input_dim |
| 295 | |
| 296 | |
| 297 | def patchify(self, hidden_states): |
no test coverage detected