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. :return: an [N x K] Tensor of outputs.
(self, x, timesteps)
| 1514 | self.middle_block.apply(convert_module_to_f32) |
| 1515 | |
| 1516 | def forward(self, x, timesteps): |
| 1517 | """ |
| 1518 | Apply the model to an input batch. |
| 1519 | :param x: an [N x C x ...] Tensor of inputs. |
| 1520 | :param timesteps: a 1-D batch of timesteps. |
| 1521 | :return: an [N x K] Tensor of outputs. |
| 1522 | """ |
| 1523 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) |
| 1524 | |
| 1525 | result_list = [] |
| 1526 | results = {} |
| 1527 | h = x.type(self.dtype) |
| 1528 | for module in self.input_blocks: |
| 1529 | last_h = h |
| 1530 | h = module(h, emb) |
| 1531 | if h.size(-1) != last_h.size(-1): |
| 1532 | result_list.append(last_h) |
| 1533 | h = self.middle_block(h, emb) |
| 1534 | result_list.append(h) |
| 1535 | |
| 1536 | assert len(result_list) == len(self.fea_tran) |
| 1537 | |
| 1538 | for i in range(len(result_list)): |
| 1539 | results[str(result_list[i].size(-1))] = self.fea_tran[i](result_list[i], emb) |
| 1540 | |
| 1541 | return results |
nothing calls this directly
no test coverage detected