| 1281 | |
| 1282 | |
| 1283 | class _WrappedModel: |
| 1284 | |
| 1285 | def __init__(self, model, timestep_map, original_num_steps): |
| 1286 | self.model = model |
| 1287 | self.timestep_map = timestep_map |
| 1288 | self.original_num_steps = original_num_steps |
| 1289 | |
| 1290 | def __call__(self, x, ts, **kwargs): |
| 1291 | map_tensor = th.tensor(self.timestep_map, |
| 1292 | device=ts.device, |
| 1293 | dtype=ts.dtype) |
| 1294 | new_ts = map_tensor[ts] |
| 1295 | return self.model(x, new_ts, **kwargs) |