(self, global_pool=False)
| 42 | |
| 43 | class SDXLControlNetUnion(torch.nn.Module): |
| 44 | def __init__(self, global_pool=False): |
| 45 | super().__init__() |
| 46 | self.time_proj = Timesteps(320) |
| 47 | self.time_embedding = torch.nn.Sequential( |
| 48 | torch.nn.Linear(320, 1280), |
| 49 | torch.nn.SiLU(), |
| 50 | torch.nn.Linear(1280, 1280) |
| 51 | ) |
| 52 | self.add_time_proj = Timesteps(256) |
| 53 | self.add_time_embedding = torch.nn.Sequential( |
| 54 | torch.nn.Linear(2816, 1280), |
| 55 | torch.nn.SiLU(), |
| 56 | torch.nn.Linear(1280, 1280) |
| 57 | ) |
| 58 | self.control_type_proj = Timesteps(256) |
| 59 | self.control_type_embedding = torch.nn.Sequential( |
| 60 | torch.nn.Linear(256 * 8, 1280), |
| 61 | torch.nn.SiLU(), |
| 62 | torch.nn.Linear(1280, 1280) |
| 63 | ) |
| 64 | self.conv_in = torch.nn.Conv2d(4, 320, kernel_size=3, padding=1) |
| 65 | |
| 66 | self.controlnet_conv_in = ControlNetConditioningLayer(channels=(3, 16, 32, 96, 256, 320)) |
| 67 | self.controlnet_transformer = ResidualAttentionBlock(320, 8) |
| 68 | self.task_embedding = torch.nn.Parameter(torch.randn(8, 320)) |
| 69 | self.spatial_ch_projs = torch.nn.Linear(320, 320) |
| 70 | |
| 71 | self.blocks = torch.nn.ModuleList([ |
| 72 | # DownBlock2D |
| 73 | ResnetBlock(320, 320, 1280), |
| 74 | PushBlock(), |
| 75 | ResnetBlock(320, 320, 1280), |
| 76 | PushBlock(), |
| 77 | DownSampler(320), |
| 78 | PushBlock(), |
| 79 | # CrossAttnDownBlock2D |
| 80 | ResnetBlock(320, 640, 1280), |
| 81 | AttentionBlock(10, 64, 640, 2, 2048), |
| 82 | PushBlock(), |
| 83 | ResnetBlock(640, 640, 1280), |
| 84 | AttentionBlock(10, 64, 640, 2, 2048), |
| 85 | PushBlock(), |
| 86 | DownSampler(640), |
| 87 | PushBlock(), |
| 88 | # CrossAttnDownBlock2D |
| 89 | ResnetBlock(640, 1280, 1280), |
| 90 | AttentionBlock(20, 64, 1280, 10, 2048), |
| 91 | PushBlock(), |
| 92 | ResnetBlock(1280, 1280, 1280), |
| 93 | AttentionBlock(20, 64, 1280, 10, 2048), |
| 94 | PushBlock(), |
| 95 | # UNetMidBlock2DCrossAttn |
| 96 | ResnetBlock(1280, 1280, 1280), |
| 97 | AttentionBlock(20, 64, 1280, 10, 2048), |
| 98 | ResnetBlock(1280, 1280, 1280), |
| 99 | PushBlock() |
| 100 | ]) |
| 101 |
no test coverage detected