|
23 | 23 | import torch |
24 | 24 |
|
25 | 25 | from ..configuration_utils import ConfigMixin, register_to_config |
26 | | -from ..utils import BaseOutput, deprecate |
| 26 | +from ..utils import BaseOutput |
27 | 27 | from .scheduling_utils import SchedulerMixin |
28 | 28 |
|
29 | 29 |
|
@@ -175,26 +175,21 @@ def _get_variance(self, timestep, prev_timestep): |
175 | 175 |
|
176 | 176 | return variance |
177 | 177 |
|
178 | | - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None, **kwargs): |
| 178 | + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): |
179 | 179 | """ |
180 | 180 | Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. |
181 | 181 |
|
182 | 182 | Args: |
183 | 183 | num_inference_steps (`int`): |
184 | 184 | the number of diffusion steps used when generating samples with a pre-trained model. |
185 | 185 | """ |
186 | | - deprecated_offset = deprecate( |
187 | | - "offset", "0.7.0", "Please pass `steps_offset` to `__init__` instead.", take_from=kwargs |
188 | | - ) |
189 | | - offset = deprecated_offset or self.config.steps_offset |
190 | | - |
191 | 186 | self.num_inference_steps = num_inference_steps |
192 | 187 | step_ratio = self.config.num_train_timesteps // self.num_inference_steps |
193 | 188 | # creates integer timesteps by multiplying by ratio |
194 | 189 | # casting to int to avoid issues when num_inference_step is power of 3 |
195 | 190 | timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) |
196 | 191 | self.timesteps = torch.from_numpy(timesteps).to(device) |
197 | | - self.timesteps += offset |
| 192 | + self.timesteps += self.config.steps_offset |
198 | 193 |
|
199 | 194 | def step( |
200 | 195 | self, |
|
0 commit comments