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)
| 967 | self.output_blocks.apply(convert_module_to_f32) |
| 968 | |
| 969 | def forward(self, x, timesteps=None, context=None, y=None, **kwargs): |
| 970 | """ |
| 971 | Apply the model to an input batch. |
| 972 | :param x: an [N x C x ...] Tensor of inputs. |
| 973 | :param timesteps: a 1-D batch of timesteps. |
| 974 | :param context: conditioning plugged in via crossattn |
| 975 | :param y: an [N] Tensor of labels, if class-conditional. |
| 976 | :return: an [N x C x ...] Tensor of outputs. |
| 977 | """ |
| 978 | assert (y is not None) == ( |
| 979 | self.num_classes is not None |
| 980 | ), "must specify y if and only if the model is class-conditional" |
| 981 | hs = [] |
| 982 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False, dtype=self.dtype) |
| 983 | emb = self.time_embed(t_emb) |
| 984 | |
| 985 | if self.num_classes is not None: |
| 986 | assert y.shape[0] == x.shape[0] |
| 987 | emb = emb + self.label_emb(y) |
| 988 | |
| 989 | # h = x.type(self.dtype) |
| 990 | h = x |
| 991 | for module in self.input_blocks: |
| 992 | h = module(h, emb, context) |
| 993 | hs.append(h) |
| 994 | h = self.middle_block(h, emb, context) |
| 995 | for module in self.output_blocks: |
| 996 | h = th.cat([h, hs.pop()], dim=1) |
| 997 | h = module(h, emb, context) |
| 998 | h = h.type(x.dtype) |
| 999 | if self.predict_codebook_ids: |
| 1000 | assert False, "not supported anymore. what the f*** are you doing?" |
| 1001 | else: |
| 1002 | return self.out(h) |
| 1003 | |
| 1004 | |
| 1005 | class NoTimeUNetModel(UNetModel): |
nothing calls this directly
no test coverage detected