| 1044 | |
| 1045 | |
| 1046 | class _WrappedModel: |
| 1047 | def __init__(self, model, timestep_map, original_num_steps): |
| 1048 | self.model = model |
| 1049 | self.timestep_map = timestep_map |
| 1050 | self.original_num_steps = original_num_steps |
| 1051 | |
| 1052 | def __call__(self, x, ts, **kwargs): |
| 1053 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) |
| 1054 | new_ts = map_tensor[ts] |
| 1055 | return self.model(x, new_ts, **kwargs) |
| 1056 | |
| 1057 | |
| 1058 | def _extract_into_tensor(arr, timesteps, broadcast_shape): |