| 34 | return len(step_sigmas) - 1 |
| 35 | |
| 36 | def forward( |
| 37 | self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict |
| 38 | ): |
| 39 | model = extra_options["model"] |
| 40 | step_sigmas = extra_options["sigmas"] |
| 41 | step = self.find_step(sigma, step_sigmas) |
| 42 | # In order to apply power multiple times, this is the same as applying power number of times equal to step |
| 43 | power = self.power**step |
| 44 | denoise_mask = denoise_mask.clone() |
| 45 | if self.only_first_frame: |
| 46 | num_channels = model.model_patcher.model.diffusion_model.in_channels |
| 47 | denoise_mask[:, :num_channels, :1] **= power |
| 48 | else: |
| 49 | denoise_mask **= power |
| 50 | # make sure to update the denoise mask in the model, to get correct timestep values for all tokens |
| 51 | for k in model.conds: |
| 52 | if "positive" in k or "negative" in k: |
| 53 | for cond in model.conds[k]: |
| 54 | if "model_conds" in cond and "denoise_mask" in cond["model_conds"]: |
| 55 | cond["model_conds"]["denoise_mask"].cond = denoise_mask |
| 56 | # print(f"DynamicConditioning: power: {power}, step: {step}, sigma: {sigma}, step_sigmas: {step_sigmas}") |
| 57 | return denoise_mask |