Apply the model to an input batch. :param x: an [N x C x ...] Tensor of inputs. :param t: a 1-D batch of timesteps. :param context: conditioning plugged in via crossattn :param y: an [N] Tensor of labels, if class-conditional. :return: an [N x C x ...
(self, x, t=None, context=None, context_ca=None, y=None, **kwargs)
| 819 | self.output_blocks.apply(convert_module_to_f32) |
| 820 | |
| 821 | def forward(self, x, t=None, context=None, context_ca=None, y=None, **kwargs): |
| 822 | """ |
| 823 | Apply the model to an input batch. |
| 824 | :param x: an [N x C x ...] Tensor of inputs. |
| 825 | :param t: a 1-D batch of timesteps. |
| 826 | :param context: conditioning plugged in via crossattn |
| 827 | :param y: an [N] Tensor of labels, if class-conditional. |
| 828 | :return: an [N x C x ...] Tensor of outputs. |
| 829 | """ |
| 830 | assert (y is not None) == ( |
| 831 | self.num_classes is not None |
| 832 | ), "must specify y if and only if the model is class-conditional" |
| 833 | hs = [] |
| 834 | t_emb = timestep_embedding(t, self.model_channels, repeat_only=False) |
| 835 | emb = self.time_embed(t_emb) |
| 836 | |
| 837 | if self.num_classes is not None: |
| 838 | assert y.shape[0] == x.shape[0] |
| 839 | emb = emb + self.label_emb(y) |
| 840 | |
| 841 | h = x.type(self.dtype) |
| 842 | if context is not None and self.concat_context: |
| 843 | h = th.cat([h, context], dim=1) |
| 844 | for module in self.input_blocks: |
| 845 | h = module(h, emb, context_ca) |
| 846 | hs.append(h) |
| 847 | h = self.middle_block(h, emb, context_ca) |
| 848 | for module in self.output_blocks: |
| 849 | h = th.cat([h, hs.pop()], dim=1) |
| 850 | h = module(h, emb, context_ca) |
| 851 | h = h.type(x.dtype) |
| 852 | if self.predict_codebook_ids: |
| 853 | return self.id_predictor(h) |
| 854 | else: |
| 855 | return self.out(h) |
| 856 | |
| 857 | def get_midblock_features(self, x, t=None, context=None, context_ca=None, y=None, **kwargs): |
| 858 | """ |
nothing calls this directly
no test coverage detected