(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
cond_key=None, return_original_cond=False, bs=None, return_x=False)
| 773 | |
| 774 | @torch.no_grad() |
| 775 | def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False, |
| 776 | cond_key=None, return_original_cond=False, bs=None, return_x=False): |
| 777 | x = super().get_input(batch, k) |
| 778 | if bs is not None: |
| 779 | x = x[:bs] |
| 780 | x = x.to(self.device) |
| 781 | encoder_posterior = self.encode_first_stage(x) |
| 782 | z = self.get_first_stage_encoding(encoder_posterior).detach() |
| 783 | |
| 784 | if self.model.conditioning_key is not None and not self.force_null_conditioning: |
| 785 | if cond_key is None: |
| 786 | cond_key = self.cond_stage_key |
| 787 | if cond_key != self.first_stage_key: |
| 788 | if cond_key in ['caption', 'coordinates_bbox', "txt"]: |
| 789 | xc = batch[cond_key] |
| 790 | elif cond_key in ['class_label', 'cls']: |
| 791 | xc = batch |
| 792 | else: |
| 793 | xc = super().get_input(batch, cond_key).to(self.device) |
| 794 | else: |
| 795 | xc = x |
| 796 | if not self.cond_stage_trainable or force_c_encode: |
| 797 | if isinstance(xc, dict) or isinstance(xc, list): |
| 798 | c = self.get_learned_conditioning(xc) |
| 799 | else: |
| 800 | c = self.get_learned_conditioning(xc.to(self.device)) |
| 801 | else: |
| 802 | c = xc |
| 803 | if bs is not None: |
| 804 | c = c[:bs] |
| 805 | |
| 806 | if self.use_positional_encodings: |
| 807 | pos_x, pos_y = self.compute_latent_shifts(batch) |
| 808 | ckey = __conditioning_keys__[self.model.conditioning_key] |
| 809 | c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y} |
| 810 | |
| 811 | else: |
| 812 | c = None |
| 813 | xc = None |
| 814 | if self.use_positional_encodings: |
| 815 | pos_x, pos_y = self.compute_latent_shifts(batch) |
| 816 | c = {'pos_x': pos_x, 'pos_y': pos_y} |
| 817 | out = [z, c] |
| 818 | if return_first_stage_outputs: |
| 819 | xrec = self.decode_first_stage(z) |
| 820 | out.extend([x, xrec]) |
| 821 | if return_x: |
| 822 | out.extend([x]) |
| 823 | if return_original_cond: |
| 824 | out.append(xc) |
| 825 | return out |
| 826 | |
| 827 | @torch.no_grad() |
| 828 | def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): |
no test coverage detected