Apply the model to an input batch and return the features from the middle block. :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
(self, x, t=None, context=None, context_ca=None, y=None, **kwargs)
| 855 | return self.out(h) |
| 856 | |
| 857 | def get_midblock_features(self, x, t=None, context=None, context_ca=None, y=None, **kwargs): |
| 858 | """ |
| 859 | Apply the model to an input batch and return the features from the middle block. |
| 860 | :param x: an [N x C x ...] Tensor of inputs. |
| 861 | :param t: a 1-D batch of timesteps. |
| 862 | :param context: conditioning plugged in via crossattn |
| 863 | :param y: an [N] Tensor of labels, if class-conditional |
| 864 | """ |
| 865 | assert (y is not None) == ( |
| 866 | self.num_classes is not None |
| 867 | ), "must specify y if and only if the model is class-conditional" |
| 868 | hs = [] |
| 869 | t_emb = timestep_embedding(t, self.model_channels, repeat_only=False) |
| 870 | emb = self.time_embed(t_emb) |
| 871 | |
| 872 | if self.num_classes is not None: |
| 873 | assert y.shape[0] == x.shape[0] |
| 874 | emb = emb + self.label_emb(y) |
| 875 | |
| 876 | h = x.type(self.dtype) |
| 877 | if context is not None and self.concat_context: |
| 878 | h = th.cat([h, context], dim=1) |
| 879 | for module in self.input_blocks: |
| 880 | h = module(h, emb, context_ca) |
| 881 | hs.append(h) |
| 882 | h = self.middle_block(h, emb, context_ca) |
| 883 | return h |
| 884 | |
| 885 | |
| 886 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected