| 187 | |
| 188 | @dataclass |
| 189 | class DPMState: |
| 190 | order: int |
| 191 | model_outputs: List[torch.Tensor] = None |
| 192 | lower_order_nums = 0 |
| 193 | |
| 194 | def __post_init__(self): |
| 195 | self.model_outputs = [None] * self.order |
| 196 | |
| 197 | def update(self, model_output: torch.Tensor): |
| 198 | for i in range(self.order - 1): |
| 199 | self.model_outputs[i] = self.model_outputs[i + 1] |
| 200 | self.model_outputs[-1] = model_output |
| 201 | |
| 202 | def update_lower_order(self): |
| 203 | if self.lower_order_nums < self.order: |
| 204 | self.lower_order_nums += 1 |
| 205 | |
| 206 | |
| 207 | def dpm_step( |