(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs)
| 25 | |
| 26 | class ControlledUnetModel(UNetModel): |
| 27 | def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs): |
| 28 | hs = [] |
| 29 | with torch.no_grad(): |
| 30 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) |
| 31 | emb = self.time_embed(t_emb) |
| 32 | h = x.type(self.dtype) |
| 33 | for module in self.input_blocks: |
| 34 | h = module(h, emb, context) |
| 35 | hs.append(h) |
| 36 | h = self.middle_block(h, emb, context) |
| 37 | |
| 38 | if control is not None: |
| 39 | h += control.pop() |
| 40 | |
| 41 | for i, module in enumerate(self.output_blocks): |
| 42 | if only_mid_control or control is None: |
| 43 | h = torch.cat([h, hs.pop()], dim=1) |
| 44 | else: |
| 45 | h = torch.cat([h, hs.pop() + control.pop()], dim=1) |
| 46 | h = module(h, emb, context) |
| 47 | |
| 48 | h = h.type(x.dtype) |
| 49 | return self.out(h) |
| 50 | |
| 51 | |
| 52 | class ControlNet(nn.Module): |
nothing calls this directly
no test coverage detected