(self, x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep=None, idx=None, scale=None, scale_emb=None, **additional_model_inputs)
| 527 | return x, s_in, alpha_cumprod_sqrt, num_sigmas, cond, uc, timesteps |
| 528 | |
| 529 | def denoise(self, x, denoiser, alpha_cumprod_sqrt, cond, uc, timestep=None, idx=None, scale=None, scale_emb=None, **additional_model_inputs): |
| 530 | |
| 531 | additional_model_inputs["is_sampling"] = True # * Used to indicate sampling |
| 532 | |
| 533 | if 'cond_inds' not in additional_model_inputs and self.cond_inds is not None: |
| 534 | additional_model_inputs['cond_inds'] = self.cond_inds |
| 535 | |
| 536 | aug_t_chunk_sampling = additional_model_inputs.get('aug_t_chunk', 0) |
| 537 | # * scale: None |
| 538 | # * We don't need to pass the scale as it will be initialized in the guider |
| 539 | |
| 540 | if not isinstance(scale, torch.Tensor) and scale == 1: |
| 541 | additional_model_inputs["idx"] = x.new_ones([x.shape[0]]) * timestep |
| 542 | |
| 543 | if self.apply_cond_aug == 'V2': |
| 544 | additional_model_inputs["aug_t_chunk"] = x.new_ones([x.shape[0]]) * aug_t_chunk_sampling |
| 545 | if scale_emb is not None: |
| 546 | additional_model_inputs["scale_emb"] = scale_emb |
| 547 | denoised = denoiser(x, alpha_cumprod_sqrt, cond, **additional_model_inputs).to(torch.float32) |
| 548 | else: |
| 549 | additional_model_inputs["idx"] = torch.cat([x.new_ones([x.shape[0]]) * timestep] * 2) |
| 550 | if self.apply_cond_aug == 'V2': |
| 551 | additional_model_inputs["aug_t_chunk"] = torch.cat([x.new_ones([x.shape[0]])] * 2) * aug_t_chunk_sampling |
| 552 | denoised = denoiser( |
| 553 | *self.guider.prepare_inputs(x, alpha_cumprod_sqrt, cond, uc), **additional_model_inputs |
| 554 | ).to(torch.float32) |
| 555 | if isinstance(self.guider, DynamicCFG) or isinstance(self.guider, TriangleCFG): |
| 556 | |
| 557 | # idx: 25, 24, 23, ..., 0 |
| 558 | if self.directly_use_idx_as_timestep: |
| 559 | sampling_step = idx |
| 560 | # sampling_step: [25, 24, 23, ...., 0] |
| 561 | |
| 562 | elif self.fix_timestep_bug: |
| 563 | sampling_step = self.num_steps - idx |
| 564 | # sampling_step: [0, 1, 2, ......] |
| 565 | |
| 566 | else: |
| 567 | sampling_step = self.num_steps - timestep |
| 568 | # sampling_step: [-974, -934, -894, ..... ,0] |
| 569 | |
| 570 | # * In the beginning of sampling |
| 571 | # timesteps: 999 |
| 572 | # self.num_steps: 25 |
| 573 | denoised = self.guider( |
| 574 | denoised, (1 - alpha_cumprod_sqrt**2) ** 0.5, step_index=sampling_step, scale=scale |
| 575 | ) |
| 576 | # step_index: -974, -934, -894 |
| 577 | else: |
| 578 | denoised = self.guider(denoised, (1 - alpha_cumprod_sqrt**2) ** 0.5, scale=scale) |
| 579 | return denoised |
| 580 | |
| 581 | def sampler_step( |
| 582 | self, |
no test coverage detected