| 36 | |
| 37 | |
| 38 | class SparseControlNet(ControlNetCLDM): |
| 39 | def __init__(self, *args,**kwargs): |
| 40 | super().__init__(*args, **kwargs) |
| 41 | hint_channels = kwargs.get("hint_channels") |
| 42 | operations: disable_weight_init_clean_groupnorm = kwargs.get("operations", disable_weight_init_clean_groupnorm) |
| 43 | device = kwargs.get("device", None) |
| 44 | self.use_simplified_conditioning_embedding = kwargs.get("use_simplified_conditioning_embedding", False) |
| 45 | if self.use_simplified_conditioning_embedding: |
| 46 | self.input_hint_block = TimestepEmbedSequential( |
| 47 | zero_module(operations.conv_nd(self.dims, hint_channels, self.model_channels, 3, padding=1, dtype=self.dtype, device=device)), |
| 48 | ) |
| 49 | |
| 50 | def forward(self, x: Tensor, hint: Tensor, timesteps, context, y=None, **kwargs): |
| 51 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype) |
| 52 | emb = self.time_embed(t_emb) |
| 53 | |
| 54 | # SparseCtrl sets noisy input to zeros |
| 55 | x = torch.zeros_like(x) |
| 56 | guided_hint = self.input_hint_block(hint, emb, context) |
| 57 | |
| 58 | out_output = [] |
| 59 | out_middle = [] |
| 60 | |
| 61 | hs = [] |
| 62 | if self.num_classes is not None: |
| 63 | assert y.shape[0] == x.shape[0] |
| 64 | emb = emb + self.label_emb(y) |
| 65 | |
| 66 | h = x |
| 67 | for module, zero_conv in zip(self.input_blocks, self.zero_convs): |
| 68 | if guided_hint is not None: |
| 69 | h = module(h, emb, context) |
| 70 | h += guided_hint |
| 71 | guided_hint = None |
| 72 | else: |
| 73 | h = module(h, emb, context) |
| 74 | out_output.append(zero_conv(h, emb, context)) |
| 75 | |
| 76 | h = self.middle_block(h, emb, context) |
| 77 | out_middle.append(self.middle_block_out(h, emb, context)) |
| 78 | |
| 79 | return {"middle": out_middle, "output": out_output} |
| 80 | |
| 81 | |
| 82 | def load_sparsectrl_motionmodel(ckpt_path: str, motion_data: dict[str, Tensor], ops=None) -> InterfaceAnimateDiffModel: |