Apply the model to an input batch. :param x: an [N x C x ...] Tensor of inputs. :param timesteps: 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
(self, x, timesteps=None, context=None, y=None,**kwargs)
| 754 | self.output_blocks.apply(convert_module_to_f32) |
| 755 | |
| 756 | def forward(self, x, timesteps=None, context=None, y=None,**kwargs): |
| 757 | """ |
| 758 | Apply the model to an input batch. |
| 759 | :param x: an [N x C x ...] Tensor of inputs. |
| 760 | :param timesteps: a 1-D batch of timesteps. |
| 761 | :param context: conditioning plugged in via crossattn |
| 762 | :param y: an [N] Tensor of labels, if class-conditional. |
| 763 | :return: an [N x C x ...] Tensor of outputs. |
| 764 | """ |
| 765 | assert (y is not None) == ( |
| 766 | self.num_classes is not None |
| 767 | ), "must specify y if and only if the model is class-conditional" |
| 768 | hs = [] |
| 769 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) |
| 770 | emb = self.time_embed(t_emb) |
| 771 | |
| 772 | if self.num_classes is not None: |
| 773 | assert y.shape[0] == x.shape[0] |
| 774 | emb = emb + self.label_emb(y) |
| 775 | |
| 776 | h = x.type(self.dtype) |
| 777 | for module in self.input_blocks: |
| 778 | h = module(h, emb, context) |
| 779 | hs.append(h) |
| 780 | h = self.middle_block(h, emb, context) |
| 781 | for module in self.output_blocks: |
| 782 | h = th.cat([h, hs.pop()], dim=1) |
| 783 | h = module(h, emb, context) |
| 784 | h = h.type(x.dtype) |
| 785 | if self.predict_codebook_ids: |
| 786 | return self.id_predictor(h) |
| 787 | else: |
| 788 | return self.out(h) |
nothing calls this directly
no test coverage detected