(self, disable_guidance_embedder=False, num_joint_blocks=5, num_single_blocks=10, num_mode=0, mode_dict={}, additional_input_dim=0)
| 7 | |
| 8 | class FluxControlNet(torch.nn.Module): |
| 9 | def __init__(self, disable_guidance_embedder=False, num_joint_blocks=5, num_single_blocks=10, num_mode=0, mode_dict={}, additional_input_dim=0): |
| 10 | super().__init__() |
| 11 | self.pos_embedder = RoPEEmbedding(3072, 10000, [16, 56, 56]) |
| 12 | self.time_embedder = TimestepEmbeddings(256, 3072) |
| 13 | self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072) |
| 14 | self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072)) |
| 15 | self.context_embedder = torch.nn.Linear(4096, 3072) |
| 16 | self.x_embedder = torch.nn.Linear(64, 3072) |
| 17 | |
| 18 | self.blocks = torch.nn.ModuleList([FluxJointTransformerBlock(3072, 24) for _ in range(num_joint_blocks)]) |
| 19 | self.single_blocks = torch.nn.ModuleList([FluxSingleTransformerBlock(3072, 24) for _ in range(num_single_blocks)]) |
| 20 | |
| 21 | self.controlnet_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_joint_blocks)]) |
| 22 | self.controlnet_single_blocks = torch.nn.ModuleList([torch.nn.Linear(3072, 3072) for _ in range(num_single_blocks)]) |
| 23 | |
| 24 | self.mode_dict = mode_dict |
| 25 | self.controlnet_mode_embedder = torch.nn.Embedding(num_mode, 3072) if len(mode_dict) > 0 else None |
| 26 | self.controlnet_x_embedder = torch.nn.Linear(64 + additional_input_dim, 3072) |
| 27 | |
| 28 | |
| 29 | def prepare_image_ids(self, latents): |
no test coverage detected