Draws samples from an interpolated cosine timestep distribution (from simple diffusion).
(shape, image_d, noise_d_low, noise_d_high, sigma_data=1., min_value=1e-3, max_value=1e3, device='cpu', dtype=torch.float32)
| 353 | |
| 354 | |
| 355 | def rand_cosine_interpolated(shape, image_d, noise_d_low, noise_d_high, sigma_data=1., min_value=1e-3, max_value=1e3, device='cpu', dtype=torch.float32): |
| 356 | """Draws samples from an interpolated cosine timestep distribution (from simple diffusion).""" |
| 357 | |
| 358 | def logsnr_schedule_cosine(t, logsnr_min, logsnr_max): |
| 359 | t_min = math.atan(math.exp(-0.5 * logsnr_max)) |
| 360 | t_max = math.atan(math.exp(-0.5 * logsnr_min)) |
| 361 | return -2 * torch.log(torch.tan(t_min + t * (t_max - t_min))) |
| 362 | |
| 363 | def logsnr_schedule_cosine_shifted(t, image_d, noise_d, logsnr_min, logsnr_max): |
| 364 | shift = 2 * math.log(noise_d / image_d) |
| 365 | return logsnr_schedule_cosine(t, logsnr_min - shift, logsnr_max - shift) + shift |
| 366 | |
| 367 | def logsnr_schedule_cosine_interpolated(t, image_d, noise_d_low, noise_d_high, logsnr_min, logsnr_max): |
| 368 | logsnr_low = logsnr_schedule_cosine_shifted(t, image_d, noise_d_low, logsnr_min, logsnr_max) |
| 369 | logsnr_high = logsnr_schedule_cosine_shifted(t, image_d, noise_d_high, logsnr_min, logsnr_max) |
| 370 | return torch.lerp(logsnr_low, logsnr_high, t) |
| 371 | |
| 372 | logsnr_min = -2 * math.log(min_value / sigma_data) |
| 373 | logsnr_max = -2 * math.log(max_value / sigma_data) |
| 374 | u = stratified_with_settings(shape, device=device, dtype=dtype) |
| 375 | logsnr = logsnr_schedule_cosine_interpolated(u, image_d, noise_d_low, noise_d_high, logsnr_min, logsnr_max) |
| 376 | return torch.exp(-logsnr / 2) * sigma_data |
| 377 | |
| 378 | |
| 379 | def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32): |
nothing calls this directly
no test coverage detected