(self, batch, use_ema: bool = True, **kwargs)
| 607 | # TODO: insert self.inference into validation_step and evaluate_and_visualize_batch |
| 608 | # to avoid any inconsistencies. |
| 609 | def inference(self, batch, use_ema: bool = True, **kwargs): |
| 610 | # check for precomputed latents |
| 611 | if "x0_latent" in batch: |
| 612 | x0_latent = batch["x0_latent"] |
| 613 | x0_latent = x0_latent * self.scale_factor |
| 614 | elif "x0" in batch: |
| 615 | x0 = batch["x0"] |
| 616 | x0_latent = self.encode_first_stage(x0) |
| 617 | else: |
| 618 | # noise or x0/x0_latent required to obtain shape for inference |
| 619 | assert (self.context_key != "x0_latent") & (self.context_key != "x0"), "x0_latent or x0 required for conditioning" |
| 620 | assert self.start_from_noise, "Only models starting from noise are allowed without x0 and x0_latent" |
| 621 | assert "noise" in batch, "Noise required for inference w/o x0 and x0_latent" |
| 622 | noise = batch["noise"] |
| 623 | x0_latent = noise # ignored |
| 624 | batch["x0_latent"] = x0_latent # scaled (ignored if starting from noise) |
| 625 | |
| 626 | """ cross-attention conditioning """ |
| 627 | if exists(self.cond_stage): |
| 628 | # fetch conditioning from raw batch |
| 629 | conditioning = batch[self.conditioning_key] |
| 630 | conditioning = self.cond_stage(conditioning) |
| 631 | else: |
| 632 | conditioning = None |
| 633 | |
| 634 | """ concatenated context """ |
| 635 | if exists(self.context_key): |
| 636 | # fetch context from preprocessed batch |
| 637 | context = batch[self.context_key] |
| 638 | else: |
| 639 | context = None |
| 640 | |
| 641 | """ input """ |
| 642 | # define x0 |
| 643 | if self.start_from_noise: |
| 644 | x_source = batch.get("noise", torch.randn_like(x0_latent)) |
| 645 | else: |
| 646 | x_source = x0_latent |
| 647 | |
| 648 | # noise x0 |
| 649 | if self.noise_image: |
| 650 | x_source = self.diffusion.q_sample(x_start=x_source, t=self.noising_step) |
| 651 | |
| 652 | """ prediction """ |
| 653 | if use_ema: |
| 654 | assert exists(self.ema_model), "Cannot use EMA for inference without EMA model" |
| 655 | model = self.ema_model if use_ema else self.model |
| 656 | x1_latent_pred = model.generate( |
| 657 | x_source, context=context, |
| 658 | context_ca=conditioning, |
| 659 | **kwargs |
| 660 | ) |
| 661 | |
| 662 | # decode |
| 663 | x1_pred = self.decode_first_stage(x1_latent_pred) |
| 664 | |
| 665 | return x1_pred |
| 666 |
nothing calls this directly
no test coverage detected