(self, batch, batch_idx)
| 549 | return loss |
| 550 | |
| 551 | def validation_step(self, batch, batch_idx): |
| 552 | # save samples for visualization |
| 553 | if self.validation_samples is None: |
| 554 | self.validation_samples = { |
| 555 | k: ( |
| 556 | v[:self.n_images_to_vis].clone() |
| 557 | if isinstance(v, Tensor) else v[:self.n_images_to_vis] |
| 558 | ) |
| 559 | for k, v in batch.items() |
| 560 | } |
| 561 | |
| 562 | """ extract data """ |
| 563 | data = self.extract_from_batch(batch) |
| 564 | x1, x1_latent = data["x1"], data["x1_latent"] |
| 565 | x0, x0_latent = data["x0"], data["x0_latent"] |
| 566 | |
| 567 | """ cross-attention conditioning """ |
| 568 | if exists(self.cond_stage): |
| 569 | # fetch conditioning from raw batch |
| 570 | conditioning = batch[self.conditioning_key] |
| 571 | conditioning = self.cond_stage(conditioning) |
| 572 | else: |
| 573 | conditioning = None |
| 574 | |
| 575 | """ concatenated context """ |
| 576 | if exists(self.context_key): |
| 577 | # fetch context from preprocessed batch |
| 578 | context = data[self.context_key] |
| 579 | else: |
| 580 | context = None |
| 581 | |
| 582 | """ input """ |
| 583 | # define x0 |
| 584 | if self.start_from_noise: |
| 585 | x_source = batch.get("noise", torch.randn_like(x1_latent)) |
| 586 | else: |
| 587 | x_source = x0_latent |
| 588 | |
| 589 | # noise x0 |
| 590 | if self.noise_image: |
| 591 | x_source = self.diffusion.q_sample(x_start=x_source, t=self.noising_step) |
| 592 | |
| 593 | """ prediction """ |
| 594 | model = self.ema_model if self.use_ema_for_sampling else self.model |
| 595 | sample_kwargs = dict(num_steps=self.sampling_steps) if hasattr(self, "sampling_steps") else {} |
| 596 | x1_latent_pred = model.generate(x_source, context=context, context_ca=conditioning, sample_kwargs=sample_kwargs) |
| 597 | |
| 598 | # decode |
| 599 | x1_pred = self.decode_first_stage(x1_latent_pred) |
| 600 | |
| 601 | """ metrics """ |
| 602 | self.metric_tracker(x1, x1_pred) |
| 603 | |
| 604 | if self.stop_training: |
| 605 | self.stop_training_method() |
| 606 | |
| 607 | # TODO: insert self.inference into validation_step and evaluate_and_visualize_batch |
| 608 | # to avoid any inconsistencies. |
nothing calls this directly
no test coverage detected