Coverage for wrapper/vibevoice/schedule/dpm

1# mypy: ignore-errors

2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/schedule/dpm_solver.py

6# Licensed under the Apache License, Version 2.0 (the "License");

7# you may not use this file except in compliance with the License.

8# You may obtain a copy of the License at

10# http://www.apache.org/licenses/LICENSE-2.0

11#

12# Unless required by applicable law or agreed to in writing, software

13# distributed under the License is distributed on an "AS IS" BASIS,

14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

15# See the License for the specific language governing permissions and

16# limitations under the License.

18# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver

20import math

21from typing import Any, List, Optional, Tuple, Union

23import numpy as np

24import torch

26from diffusers.configuration_utils import ConfigMixin, register_to_config

27from diffusers.utils import deprecate

28from diffusers.utils.torch_utils import randn_tensor

29from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput

32def betas_for_alpha_bar(

33 num_diffusion_timesteps: int,

34 max_beta: float = 0.999,

35 alpha_transform_type: str = "cosine",

36) -> torch.Tensor:

37 """

38 Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of

39 (1-beta) over time from t = [0,1].

41 Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up

42 to that part of the diffusion process.

45 Args:

46 num_diffusion_timesteps (`int`): the number of betas to produce.

47 max_beta (`float`): the maximum beta to use; use values lower than 1 to

48 prevent singularities.

49 alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.

50 Choose from `cosine` or `exp`

52 Returns:

53 betas (`np.ndarray`): the betas used by the scheduler to step the model outputs

54 """

55 if alpha_transform_type == "cosine":

57 def alpha_bar_fn(t: float) -> float:

58 return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2

59 # return math.cos(t * math.pi / 2 * 0.95) ** 2

61 elif alpha_transform_type == "exp":

63 def alpha_bar_fn(t: float) -> float:

64 return math.exp(t * -12.0)

66 elif alpha_transform_type == "cauchy":

67 # µ + γ tan (π (0.5 - x)) γ = 1, µ = 3

68 # alpha^2 = 1-1/(exp(λ)+1)

69 _gamma: float = 1.0

70 _mu_c: float = 3.0

72 def alpha_bar_fn(t: float) -> float:

73 snr = _mu_c + _gamma * math.tan(math.pi * (0.5 - t) * 0.9)

74 return 1 - 1 / (math.exp(snr) + 1.1)

76 elif alpha_transform_type == "laplace":

77 # µ − bsgn(0.5 − t) log(1 − 2|t − 0.5|) µ = 0, b = 1

78 _mu_l: float = 0.0

79 _b: float = 1.0

81 def alpha_bar_fn(t: float) -> float:

82 snr = _mu_l - _b * math.copysign(1, 0.5 - t) * math.log(1 - 2 * abs(t - 0.5) * 0.98)

83 return 1 - 1 / (math.exp(snr) + 1.02)

85 else:

86 raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")

88 betas = []

89 for i in range(num_diffusion_timesteps):

90 t1 = i / num_diffusion_timesteps

91 t2 = (i + 1) / num_diffusion_timesteps

92 betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))

93 return torch.tensor(betas, dtype=torch.float32)

96# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr

97def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:

98 """

99 Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)

100

101

102 Args:

103 betas (`torch.Tensor`):

104 the betas that the scheduler is being initialized with.

105

106 Returns:

107 `torch.Tensor`: rescaled betas with zero terminal SNR

108 """

109 # Convert betas to alphas_bar_sqrt

110 alphas = 1.0 - betas

111 alphas_cumprod = torch.cumprod(alphas, dim=0)

112 alphas_bar_sqrt = alphas_cumprod.sqrt()

113

114 # Store old values.

115 alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()

116 alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()

117

118 # Shift so the last timestep is zero.

119 alphas_bar_sqrt -= alphas_bar_sqrt_T

120

121 # Scale so the first timestep is back to the old value.

122 alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)

123

124 # Convert alphas_bar_sqrt to betas

125 alphas_bar = alphas_bar_sqrt**2 # Revert sqrt

126 alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod

127 alphas = torch.cat([alphas_bar[0:1], alphas])

128 betas = 1 - alphas

129

130 return betas

131

132

133class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):

134 """

135 `DPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.

136

137 This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic

138 methods the library implements for all schedulers such as loading and saving.

139

140 Args:

141 num_train_timesteps (`int`, defaults to 1000):

142 The number of diffusion steps to train the model.

143 beta_start (`float`, defaults to 0.0001):

144 The starting `beta` value of inference.

145 beta_end (`float`, defaults to 0.02):

146 The final `beta` value.

147 beta_schedule (`str`, defaults to `"linear"`):

148 The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from

149 `linear`, `scaled_linear`, or `squaredcos_cap_v2`.

150 trained_betas (`np.ndarray`, *optional*):

151 Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.

152 solver_order (`int`, defaults to 2):

153 The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided

154 sampling, and `solver_order=3` for unconditional sampling.

155 prediction_type (`str`, defaults to `epsilon`, *optional*):

156 Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),

157 `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen

158 Video](https://imagen.research.google/video/paper.pdf) paper).

159 thresholding (`bool`, defaults to `False`):

160 Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such

161 as Stable Diffusion.

162 dynamic_thresholding_ratio (`float`, defaults to 0.995):

163 The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.

164 sample_max_value (`float`, defaults to 1.0):

165 The threshold value for dynamic thresholding. Valid only when `thresholding=True` and

166 `algorithm_type="dpmsolver++"`.

167 algorithm_type (`str`, defaults to `dpmsolver++`):

168 Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The

169 `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)

170 paper, and the `dpmsolver++` type implements the algorithms in the

171 [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or

172 `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.

173 solver_type (`str`, defaults to `midpoint`):

174 Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the

175 sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.

176 lower_order_final (`bool`, defaults to `True`):

177 Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can

178 stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.

179 euler_at_final (`bool`, defaults to `False`):

180 Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail

181 richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference

182 steps, but sometimes may result in blurring.

183 use_karras_sigmas (`bool`, *optional*, defaults to `False`):

184 Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,

185 the sigmas are determined according to a sequence of noise levels {σi}.

186 use_lu_lambdas (`bool`, *optional*, defaults to `False`):

187 Whether to use the uniform-logSNR for step sizes proposed by Lu's DPM-Solver in the noise schedule during

188 the sampling process. If `True`, the sigmas and time steps are determined according to a sequence of

189 `lambda(t)`.

190 final_sigmas_type (`str`, defaults to `"zero"`):

191 The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final

192 sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.

193 lambda_min_clipped (`float`, defaults to `-inf`):

194 Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the

195 cosine (`squaredcos_cap_v2`) noise schedule.

196 variance_type (`str`, *optional*):

197 Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output

198 contains the predicted Gaussian variance.

199 timestep_spacing (`str`, defaults to `"linspace"`):

200 The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and

201 Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.

202 steps_offset (`int`, defaults to 0):

203 An offset added to the inference steps, as required by some model families.

204 rescale_betas_zero_snr (`bool`, defaults to `False`):

205 Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and

206 dark samples instead of limiting it to samples with medium brightness. Loosely related to

207 [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).

208 """

209

210 _compatibles = [e.name for e in KarrasDiffusionSchedulers]

211 order = 1

212

213 @register_to_config

214 def __init__(

215 self,

216 num_train_timesteps: int = 1000,

217 beta_start: float = 0.0001,

218 beta_end: float = 0.02,

219 beta_schedule: str = "linear",

220 trained_betas: Optional[Union[np.ndarray, List[float]]] = None,

221 solver_order: int = 2,

222 prediction_type: str = "epsilon",

223 thresholding: bool = False,

224 dynamic_thresholding_ratio: float = 0.995,

225 sample_max_value: float = 1.0,

226 algorithm_type: str = "dpmsolver++",

227 solver_type: str = "midpoint",

228 lower_order_final: bool = True,

229 euler_at_final: bool = False,

230 use_karras_sigmas: Optional[bool] = False,

231 use_lu_lambdas: Optional[bool] = False,

232 final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"

233 lambda_min_clipped: float = -float("inf"),

234 variance_type: Optional[str] = None,

235 timestep_spacing: str = "linspace",

236 steps_offset: int = 0,

237 rescale_betas_zero_snr: bool = False,

238 ) -> None:

239 if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:

240 deprecation_message = (

241 f"algorithm_type {algorithm_type} is deprecated and will be removed "

242 "in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"

243 )

244 deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)

245

246 if trained_betas is not None:

247 self.betas = torch.tensor(trained_betas, dtype=torch.float32)

248 elif beta_schedule == "linear":

249 self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)

250 elif beta_schedule == "scaled_linear":

251 # this schedule is very specific to the latent diffusion model.

252 self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2

253 elif beta_schedule == "squaredcos_cap_v2" or beta_schedule == "cosine":

254 # Glide cosine schedule

255 self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cosine")

256 elif beta_schedule == "cauchy":

257 self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cauchy")

258 elif beta_schedule == "laplace":

259 self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="laplace")

260 else:

261 raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")

262

263 if rescale_betas_zero_snr:

264 self.betas = rescale_zero_terminal_snr(self.betas)

265

266 self.alphas = 1.0 - self.betas

267 self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

268

269 if rescale_betas_zero_snr:

270 # Close to 0 without being 0 so first sigma is not inf

271 # FP16 smallest positive subnormal works well here

272 self.alphas_cumprod[-1] = 2**-24

273

274 # Currently we only support VP-type noise schedule

275 self.alpha_t = torch.sqrt(self.alphas_cumprod)

276 self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)

277 self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)

278 self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5

279

280 # standard deviation of the initial noise distribution

281 self.init_noise_sigma = 1.0

282

283 # settings for DPM-Solver

284 if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:

285 if algorithm_type == "deis":

286 self.register_to_config(algorithm_type="dpmsolver++")

287 else:

288 raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}")

289

290 if solver_type not in ["midpoint", "heun"]:

291 if solver_type in ["logrho", "bh1", "bh2"]:

292 self.register_to_config(solver_type="midpoint")

293 else:

294 raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")

295

296 if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero":

297 raise ValueError(

298 f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. "

299 "Please choose `sigma_min` instead."

300 )

301

302 # setable values

303 self.num_inference_steps: Optional[int] = None

304 timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()

305 self.timesteps = torch.from_numpy(timesteps)

306 self.model_outputs = [None] * solver_order

307 self.lower_order_nums = 0

308 self._step_index: Optional[int] = None

309 self._begin_index: Optional[int] = None

310 self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication

311

312 @property

313 def step_index(self) -> Optional[int]:

314 """

315 The index counter for current timestep. It will increase 1 after each scheduler step.

316 """

317 return self._step_index

318

319 @property

320 def begin_index(self) -> Optional[int]:

321 """

322 The index for the first timestep. It should be set from pipeline with `set_begin_index` method.

323 """

324 return self._begin_index

325

326 def set_begin_index(self, begin_index: int = 0) -> None:

327 """

328 Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

329

330 Args:

331 begin_index (`int`):

332 The begin index for the scheduler.

333 """

334 self._begin_index = begin_index

335

336 def set_timesteps(

337 self,

338 num_inference_steps: Optional[int] = None,

339 device: Union[str, torch.device] = None,

340 timesteps: Optional[List[int]] = None,

341 ) -> None:

342 """

343 Sets the discrete timesteps used for the diffusion chain (to be run before inference).

344

345 Args:

346 num_inference_steps (`int`):

347 The number of diffusion steps used when generating samples with a pre-trained model.

348 device (`str` or `torch.device`, *optional*):

349 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.

350 timesteps (`List[int]`, *optional*):

351 Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated

352 based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`

353 must be `None`, and `timestep_spacing` attribute will be ignored.

354 """

355 if num_inference_steps is None and timesteps is None:

356 raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.")

357 if num_inference_steps is not None and timesteps is not None:

358 raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")

359 if timesteps is not None and self.config.use_karras_sigmas:

360 raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`")

361 if timesteps is not None and self.config.use_lu_lambdas:

362 raise ValueError("Cannot use `timesteps` with `config.use_lu_lambdas = True`")

363

364 if timesteps is not None:

365 ts: np.ndarray = np.array(timesteps).astype(np.int64)

366 else:

367 # Clipping the minimum of all lambda(t) for numerical stability.

368 # This is critical for cosine (squaredcos_cap_v2) noise schedule.

369 clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped)

370 last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item()

371

372 # "linspace", "leading", "trailing" corresponds to annotation of

373 # Table 2. of https://arxiv.org/abs/2305.08891

374 assert num_inference_steps is not None

375 if self.config.timestep_spacing == "linspace":

376 ts = (

377 np.linspace(0, last_timestep - 1, num_inference_steps + 1)

378 .round()[::-1][:-1]

379 .copy()

380 .astype(np.int64)

381 )

382 elif self.config.timestep_spacing == "leading":

383 step_ratio = last_timestep // (num_inference_steps + 1)

384 # creates integer timesteps by multiplying by ratio

385 # casting to int to avoid issues when num_inference_step is power of 3

386 _ts_leading = (

387 (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)

388 )

389 _ts_leading += self.config.steps_offset # type: ignore[attr-defined]

390 ts = _ts_leading

391 elif self.config.timestep_spacing == "trailing":

392 step_ratio = self.config.num_train_timesteps / num_inference_steps

393 # creates integer timesteps by multiplying by ratio

394 # casting to int to avoid issues when num_inference_step is power of 3

395 _ts_trailing = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64)

396 _ts_trailing -= 1

397 ts = _ts_trailing

398 else:

399 raise ValueError(

400 f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of "

401 "'linspace', 'leading' or 'trailing'."

402 )

403

404 sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)

405 log_sigmas = np.log(sigmas)

406

407 if self.config.use_karras_sigmas:

408 assert num_inference_steps is not None

409 sigmas = np.flip(sigmas).copy()

410 sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)

411 ts = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()

412 elif self.config.use_lu_lambdas:

413 assert num_inference_steps is not None

414 lambdas = np.flip(log_sigmas.copy())

415 lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps)

416 sigmas = np.exp(lambdas)

417 ts = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()

418 else:

419 sigmas = np.interp(ts, np.arange(0, len(sigmas)), sigmas)

420

421 if self.config.final_sigmas_type == "sigma_min":

422 sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5

423 elif self.config.final_sigmas_type == "zero":

424 sigma_last = 0

425 else:

426 raise ValueError(

427 f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"

428 )

429

430 sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)

431

432 self.sigmas = torch.from_numpy(sigmas)

433 self.timesteps = torch.from_numpy(ts).to(device=device, dtype=torch.int64)

434

435 self.num_inference_steps = len(ts)

436

437 self.model_outputs = [

438 None,

439 ] * self.config.solver_order

440 self.lower_order_nums = 0

441

442 # add an index counter for schedulers that allow duplicated timesteps

443 self._step_index = None

444 self._begin_index = None

445 self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication

446

447 # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample

448 def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:

449 """

450 "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the

451 prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by

452 s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing

453 pixels from saturation at each step. We find that dynamic thresholding results in significantly better

454 photorealism as well as better image-text alignment, especially when using very large guidance weights."

455

456 https://arxiv.org/abs/2205.11487

457 """

458 dtype = sample.dtype

459 batch_size, channels, *remaining_dims = sample.shape

460

461 if dtype not in (torch.float32, torch.float64):

462 sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half

463

464 # Flatten sample for doing quantile calculation along each image

465 sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))

466

467 abs_sample = sample.abs() # "a certain percentile absolute pixel value"

468

469 s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)

470 s = torch.clamp(

471 s, min=1, max=self.config.sample_max_value

472 ) # When clamped to min=1, equivalent to standard clipping to [-1, 1]

473 s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0

474 sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s"

475

476 sample = sample.reshape(batch_size, channels, *remaining_dims)

477 sample = sample.to(dtype)

478

479 return sample

480

481 # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t

482 def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray:

483 # get log sigma

484 log_sigma = np.log(np.maximum(sigma, 1e-10))

485

486 # get distribution

487 dists = log_sigma - log_sigmas[:, np.newaxis]

488

489 # get sigmas range

490 low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)

491 high_idx = low_idx + 1

492

493 low = log_sigmas[low_idx]

494 high = log_sigmas[high_idx]

495

496 # interpolate sigmas

497 w = (low - log_sigma) / (low - high)

498 w = np.clip(w, 0, 1)

499

500 # transform interpolation to time range

501 t = (1 - w) * low_idx + w * high_idx

502 t = t.reshape(sigma.shape)

503 return t

504

505 def _sigma_to_alpha_sigma_t(self, sigma: Any) -> Tuple[Any, Any]:

506 alpha_t = 1 / ((sigma**2 + 1) ** 0.5)

507 sigma_t = sigma * alpha_t

508

509 return alpha_t, sigma_t

510

511 # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras

512 def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:

513 """Constructs the noise schedule of Karras et al. (2022)."""

514

515 # Hack to make sure that other schedulers which copy this function don't break

516 # TODO: Add this logic to the other schedulers

517 if hasattr(self.config, "sigma_min"):

518 sigma_min = self.config.sigma_min

519 else:

520 sigma_min = None

521

522 if hasattr(self.config, "sigma_max"):

523 sigma_max = self.config.sigma_max

524 else:

525 sigma_max = None

526

527 sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()

528 sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()

529

530 rho = 7.0 # 7.0 is the value used in the paper

531 ramp = np.linspace(0, 1, num_inference_steps)

532 min_inv_rho = sigma_min ** (1 / rho)

533 max_inv_rho = sigma_max ** (1 / rho)

534 sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho

535 return sigmas

536

537 def _convert_to_lu(self, in_lambdas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:

538 """Constructs the noise schedule of Lu et al. (2022)."""

539

540 lambda_min: float = in_lambdas[-1].item()

541 lambda_max: float = in_lambdas[0].item()

542

543 rho = 1.0 # 1.0 is the value used in the paper

544 ramp = np.linspace(0, 1, num_inference_steps)

545 min_inv_rho = lambda_min ** (1 / rho)

546 max_inv_rho = lambda_max ** (1 / rho)

547 lambdas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho

548 return lambdas

549

550 def convert_model_output(

551 self,

552 model_output: torch.Tensor,

553 *args: Any,

554 sample: Optional[torch.Tensor] = None,

555 **kwargs: Any,

556 ) -> torch.Tensor:

557 """

558 Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is

559 designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an

560 integral of the data prediction model.

561

562 <Tip>

563

564 The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise

565 prediction and data prediction models.

566

567 </Tip>

568

569 Args:

570 model_output (`torch.Tensor`):

571 The direct output from the learned diffusion model.

572 sample (`torch.Tensor`):

573 A current instance of a sample created by the diffusion process.

574

575 Returns:

576 `torch.Tensor`:

577 The converted model output.

578 """

579 timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)

580 if sample is None:

581 if len(args) > 1:

582 sample = args[1]

583 else:

584 raise ValueError("missing `sample` as a required keyward argument")

585 if timestep is not None:

586 deprecate(

587 "timesteps",

588 "1.0.0",

589 "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via "

590 "an internal counter `self.step_index`",

591 )

592

593 # DPM-Solver++ needs to solve an integral of the data prediction model.

594 if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:

595 if self.config.prediction_type == "epsilon":

596 # DPM-Solver and DPM-Solver++ only need the "mean" output.

597 if self.config.variance_type in ["learned", "learned_range"]:

598 model_output = model_output[:, :3]

599 sigma = self.sigmas[self.step_index]

600 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

601 x0_pred = (sample - sigma_t * model_output) / alpha_t

602 elif self.config.prediction_type == "sample":

603 x0_pred = model_output

604 elif self.config.prediction_type == "v_prediction":

605 sigma = self.sigmas[self.step_index]

606 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

607 x0_pred = alpha_t * sample - sigma_t * model_output

608 else:

609 raise ValueError(

610 f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"

611 " `v_prediction` for the DPMSolverMultistepScheduler."

612 )

613

614 if self.config.thresholding:

615 x0_pred = self._threshold_sample(x0_pred)

616

617 return x0_pred

618

619 # DPM-Solver needs to solve an integral of the noise prediction model.

620 elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:

621 if self.config.prediction_type == "epsilon":

622 # DPM-Solver and DPM-Solver++ only need the "mean" output.

623 if self.config.variance_type in ["learned", "learned_range"]:

624 epsilon = model_output[:, :3]

625 else:

626 epsilon = model_output

627 elif self.config.prediction_type == "sample":

628 sigma = self.sigmas[self.step_index]

629 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

630 epsilon = (sample - alpha_t * model_output) / sigma_t

631 elif self.config.prediction_type == "v_prediction":

632 sigma = self.sigmas[self.step_index]

633 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

634 epsilon = alpha_t * model_output + sigma_t * sample

635 else:

636 raise ValueError(

637 f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"

638 " `v_prediction` for the DPMSolverMultistepScheduler."

639 )

640

641 if self.config.thresholding:

642 sigma = self.sigmas[self.step_index]

643 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)

644 x0_pred = (sample - sigma_t * epsilon) / alpha_t

645 x0_pred = self._threshold_sample(x0_pred)

646 epsilon = (sample - alpha_t * x0_pred) / sigma_t

647

648 return epsilon

649

650 def dpm_solver_first_order_update(

651 self,

652 model_output: torch.Tensor,

653 *args: Any,

654 sample: Optional[torch.Tensor] = None,

655 noise: Optional[torch.Tensor] = None,

656 **kwargs: Any,

657 ) -> torch.Tensor:

658 """

659 One step for the first-order DPMSolver (equivalent to DDIM).

660

661 Args:

662 model_output (`torch.Tensor`):

663 The direct output from the learned diffusion model.

664 sample (`torch.Tensor`):

665 A current instance of a sample created by the diffusion process.

666

667 Returns:

668 `torch.Tensor`:

669 The sample tensor at the previous timestep.

670 """

671 timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)

672 prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)

673 if sample is None:

674 if len(args) > 2:

675 sample = args[2]

676 else:

677 raise ValueError(" missing `sample` as a required keyward argument")

678 if timestep is not None:

679 deprecate(

680 "timesteps",

681 "1.0.0",

682 "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an "

683 "internal counter `self.step_index`",

684 )

685

686 if prev_timestep is not None:

687 deprecate(

688 "prev_timestep",

689 "1.0.0",

690 "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled "

691 "via an internal counter `self.step_index`",

692 )

693

694 assert self.step_index is not None

695 sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]

696 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)

697 alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)

698 lambda_t = torch.log(alpha_t) - torch.log(sigma_t)

699 lambda_s = torch.log(alpha_s) - torch.log(sigma_s)

700

701 h = lambda_t - lambda_s

702 if self.config.algorithm_type == "dpmsolver++":

703 x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output

704 elif self.config.algorithm_type == "dpmsolver":

705 x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output

706 elif self.config.algorithm_type == "sde-dpmsolver++":

707 assert noise is not None

708 x_t = (

709 (sigma_t / sigma_s * torch.exp(-h)) * sample

710 + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output

711 + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise

712 )

713 elif self.config.algorithm_type == "sde-dpmsolver":

714 assert noise is not None

715 x_t = (

716 (alpha_t / alpha_s) * sample

717 - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output

718 + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise

719 )

720 return x_t

721

722 def multistep_dpm_solver_second_order_update(

723 self,

724 model_output_list: List[torch.Tensor],

725 *args: Any,

726 sample: Optional[torch.Tensor] = None,

727 noise: Optional[torch.Tensor] = None,

728 **kwargs: Any,

729 ) -> torch.Tensor:

730 """

731 One step for the second-order multistep DPMSolver.

732

733 Args:

734 model_output_list (`List[torch.Tensor]`):

735 The direct outputs from learned diffusion model at current and latter timesteps.

736 sample (`torch.Tensor`):

737 A current instance of a sample created by the diffusion process.

738

739 Returns:

740 `torch.Tensor`:

741 The sample tensor at the previous timestep.

742 """

743 timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)

744 prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)

745 if sample is None:

746 if len(args) > 2:

747 sample = args[2]

748 else:

749 raise ValueError(" missing `sample` as a required keyward argument")

750 if timestep_list is not None:

751 deprecate(

752 "timestep_list",

753 "1.0.0",

754 "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled "

755 "via an internal counter `self.step_index`",

756 )

757

758 if prev_timestep is not None:

759 deprecate(

760 "prev_timestep",

761 "1.0.0",

762 "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled "

763 "via an internal counter `self.step_index`",

764 )

765

766 assert self.step_index is not None

767 sigma_t, sigma_s0, sigma_s1 = (

768 self.sigmas[self.step_index + 1],

769 self.sigmas[self.step_index],

770 self.sigmas[self.step_index - 1],

771 )

772

773 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)

774 alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

775 alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)

776

777 lambda_t = torch.log(alpha_t) - torch.log(sigma_t)

778 lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)

779 lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)

780

781 m0, m1 = model_output_list[-1], model_output_list[-2]

782

783 h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1

784 r0 = h_0 / h

785 D0, D1 = m0, (1.0 / r0) * (m0 - m1)

786 if self.config.algorithm_type == "dpmsolver++":

787 # See https://arxiv.org/abs/2211.01095 for detailed derivations

788 if self.config.solver_type == "midpoint":

789 x_t = (

790 (sigma_t / sigma_s0) * sample

791 - (alpha_t * (torch.exp(-h) - 1.0)) * D0

792 - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1

793 )

794 elif self.config.solver_type == "heun":

795 x_t = (

796 (sigma_t / sigma_s0) * sample

797 - (alpha_t * (torch.exp(-h) - 1.0)) * D0

798 + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1

799 )

800 elif self.config.algorithm_type == "dpmsolver":

801 # See https://arxiv.org/abs/2206.00927 for detailed derivations

802 if self.config.solver_type == "midpoint":

803 x_t = (

804 (alpha_t / alpha_s0) * sample

805 - (sigma_t * (torch.exp(h) - 1.0)) * D0

806 - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1

807 )

808 elif self.config.solver_type == "heun":

809 x_t = (

810 (alpha_t / alpha_s0) * sample

811 - (sigma_t * (torch.exp(h) - 1.0)) * D0

812 - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1

813 )

814 elif self.config.algorithm_type == "sde-dpmsolver++":

815 assert noise is not None

816 if self.config.solver_type == "midpoint":

817 x_t = (

818 (sigma_t / sigma_s0 * torch.exp(-h)) * sample

819 + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0

820 + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1

821 + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise

822 )

823 elif self.config.solver_type == "heun":

824 x_t = (

825 (sigma_t / sigma_s0 * torch.exp(-h)) * sample

826 + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0

827 + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1

828 + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise

829 )

830 elif self.config.algorithm_type == "sde-dpmsolver":

831 assert noise is not None

832 if self.config.solver_type == "midpoint":

833 x_t = (

834 (alpha_t / alpha_s0) * sample

835 - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0

836 - (sigma_t * (torch.exp(h) - 1.0)) * D1

837 + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise

838 )

839 elif self.config.solver_type == "heun":

840 x_t = (

841 (alpha_t / alpha_s0) * sample

842 - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0

843 - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1

844 + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise

845 )

846 return x_t

847

848 def multistep_dpm_solver_third_order_update(

849 self,

850 model_output_list: List[torch.Tensor],

851 *args: Any,

852 sample: Optional[torch.Tensor] = None,

853 **kwargs: Any,

854 ) -> torch.Tensor:

855 """

856 One step for the third-order multistep DPMSolver.

857

858 Args:

859 model_output_list (`List[torch.Tensor]`):

860 The direct outputs from learned diffusion model at current and latter timesteps.

861 sample (`torch.Tensor`):

862 A current instance of a sample created by diffusion process.

863

864 Returns:

865 `torch.Tensor`:

866 The sample tensor at the previous timestep.

867 """

868

869 timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)

870 prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)

871 if sample is None:

872 if len(args) > 2:

873 sample = args[2]

874 else:

875 raise ValueError(" missing`sample` as a required keyward argument")

876 if timestep_list is not None:

877 deprecate(

878 "timestep_list",

879 "1.0.0",

880 "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled "

881 "via an internal counter `self.step_index`",

882 )

883

884 if prev_timestep is not None:

885 deprecate(

886 "prev_timestep",

887 "1.0.0",

888 "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled "

889 "via an internal counter `self.step_index`",

890 )

891

892 assert self.step_index is not None

893 sigma_t, sigma_s0, sigma_s1, sigma_s2 = (

894 self.sigmas[self.step_index + 1],

895 self.sigmas[self.step_index],

896 self.sigmas[self.step_index - 1],

897 self.sigmas[self.step_index - 2],

898 )

899

900 alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)

901 alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)

902 alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)

903 alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)

904

905 lambda_t = torch.log(alpha_t) - torch.log(sigma_t)

906 lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)

907 lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)

908 lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)

909

910 m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]

911

912 h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2

913 r0, r1 = h_0 / h, h_1 / h

914 D0 = m0

915 D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)

916 D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)

917 D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)

918 if self.config.algorithm_type == "dpmsolver++":

919 # See https://arxiv.org/abs/2206.00927 for detailed derivations

920 x_t = (

921 (sigma_t / sigma_s0) * sample

922 - (alpha_t * (torch.exp(-h) - 1.0)) * D0

923 + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1

924 - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2

925 )

926 elif self.config.algorithm_type == "dpmsolver":

927 # See https://arxiv.org/abs/2206.00927 for detailed derivations

928 x_t = (

929 (alpha_t / alpha_s0) * sample

930 - (sigma_t * (torch.exp(h) - 1.0)) * D0

931 - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1

932 - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2

933 )

934 return x_t

935

936 def index_for_timestep(self, timestep: Any, schedule_timesteps: Optional[torch.Tensor] = None) -> int:

937 if schedule_timesteps is None:

938 schedule_timesteps = self.timesteps

939

940 index_candidates = (schedule_timesteps == timestep).nonzero()

941

942 if len(index_candidates) == 0:

943 step_index = len(self.timesteps) - 1

944 # The sigma index that is taken for the **very** first `step`

945 # is always the second index (or the last index if there is only 1)

946 # This way we can ensure we don't accidentally skip a sigma in

947 # case we start in the middle of the denoising schedule (e.g. for image-to-image)

948 elif len(index_candidates) > 1:

949 step_index = index_candidates[1].item()

950 else:

951 step_index = index_candidates[0].item()

952

953 return step_index

954

955 def _init_step_index(self, timestep: Any) -> None:

956 """

957 Initialize the step_index counter for the scheduler.

958 """

959

960 if self.begin_index is None:

961 if isinstance(timestep, torch.Tensor):

962 timestep = timestep.to(self.timesteps.device)

963 self._step_index = self.index_for_timestep(timestep)

964 else:

965 self._step_index = self._begin_index

966

967 def step(

968 self,

969 model_output: torch.Tensor,

970 timestep: int,

971 sample: torch.Tensor,

972 generator: Optional[torch.Generator] = None,

973 variance_noise: Optional[torch.Tensor] = None,

974 return_dict: bool = True,

975 ) -> Union[SchedulerOutput, Tuple]:

976 """

977 Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with

978 the multistep DPMSolver.

979

980 Args:

981 model_output (`torch.Tensor`):

982 The direct output from learned diffusion model.

983 timestep (`int`):

984 The current discrete timestep in the diffusion chain.

985 sample (`torch.Tensor`):

986 A current instance of a sample created by the diffusion process.

987 generator (`torch.Generator`, *optional*):

988 A random number generator.

989 variance_noise (`torch.Tensor`):

990 Alternative to generating noise with `generator` by directly providing the noise for the variance

991 itself. Useful for methods such as [`LEdits++`].

992 return_dict (`bool`):

993 Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.

994

995 Returns:

996 [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:

997 If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a

998 tuple is returned where the first element is the sample tensor.

999

1000 """

1001 if self.num_inference_steps is None:

1002 raise ValueError(

1003 "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"

1004 )

1005

1006 if self.step_index is None:

1007 self._init_step_index(timestep)

1008

1009 # Improve numerical stability for small number of steps

1010 lower_order_final = (self.step_index == len(self.timesteps) - 1) and (

1011 self.config.euler_at_final

1012 or (self.config.lower_order_final and len(self.timesteps) < 15)

1013 or self.config.final_sigmas_type == "zero"

1014 )

1015 lower_order_second = (

1016 (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15

1017 )

1018

1019 model_output = self.convert_model_output(model_output, sample=sample)

1020 for i in range(self.config.solver_order - 1):

1021 self.model_outputs[i] = self.model_outputs[i + 1]

1022 self.model_outputs[-1] = model_output

1023

1024 # Upcast to avoid precision issues when computing prev_sample

1025 sample = sample.to(torch.float32)

1026 if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:

1027 noise = randn_tensor(

1028 model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32

1029 )

1030 elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:

1031 assert variance_noise is not None

1032 noise = variance_noise.to(device=model_output.device, dtype=torch.float32)

1033 else:

1034 noise = None

1035

1036 if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:

1037 prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise)

1038 elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:

1039 prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise)

1040 else:

1041 prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample)

1042

1043 if self.lower_order_nums < self.config.solver_order:

1044 self.lower_order_nums += 1

1045

1046 # Cast sample back to expected dtype

1047 prev_sample = prev_sample.to(model_output.dtype)

1048

1049 # upon completion increase step index by one

1050 assert self._step_index is not None

1051 self._step_index += 1

1052

1053 if not return_dict:

1054 return (prev_sample,)

1055

1056 return SchedulerOutput(prev_sample=prev_sample)

1057

1058 def add_noise(

1059 self,

1060 original_samples: torch.Tensor,

1061 noise: torch.Tensor,

1062 timesteps: torch.IntTensor,

1063 ) -> torch.Tensor:

1064 # Make sure sigmas and timesteps have the same device and dtype as original_samples

1065 # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)

1066 # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)

1067 alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)

1068 sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)

1069 timesteps = timesteps.to(original_samples.device)

1070 alpha_t = alpha_t[timesteps].flatten()

1071 while len(alpha_t.shape) < len(original_samples.shape):

1072 alpha_t = alpha_t.unsqueeze(-1)

1073

1074 sigma_t = sigma_t[timesteps].flatten()

1075 while len(sigma_t.shape) < len(original_samples.shape):

1076 sigma_t = sigma_t.unsqueeze(-1)

1077 noisy_samples = alpha_t * original_samples + sigma_t * noise

1078 return noisy_samples

1079

1080 def get_velocity(

1081 self,

1082 original_samples: torch.Tensor,

1083 noise: torch.Tensor,

1084 timesteps: torch.IntTensor

1085 ) -> torch.Tensor:

1086 # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)

1087 # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)

1088 alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)

1089 sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)

1090

1091 timesteps = timesteps.to(original_samples.device)

1092 alpha_t = alpha_t[timesteps].flatten()

1093 while len(alpha_t.shape) < len(original_samples.shape):

1094 alpha_t = alpha_t.unsqueeze(-1)

1095

1096 sigma_t = sigma_t[timesteps].flatten()

1097 while len(sigma_t.shape) < len(original_samples.shape):

1098 sigma_t = sigma_t.unsqueeze(-1)

1099

1100 velocity = alpha_t * noise - sigma_t * original_samples

1101 return velocity

1102

1103 def __len__(self) -> int:

1104 return self.config.num_train_timesteps

Coverage for wrapper/vibevoice/schedule/dpm_solver.py: 48%

441 statements