From 5a93b958120f5caebdb33205073fdbc645ada8aa Mon Sep 17 00:00:00 2001 From: Roger Date: Thu, 31 Oct 2024 01:35:33 +0530 Subject: [PATCH 1/2] Refac training utils.py --- src/diffusers/training_utils.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 0e0d0ce5b568..8de9f876b4a2 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -43,6 +43,9 @@ def set_seed(seed: int): Args: seed (`int`): The seed to set. + + Returns: + `None` """ random.seed(seed) np.random.seed(seed) @@ -56,8 +59,18 @@ def set_seed(seed: int): def compute_snr(noise_scheduler, timesteps): """ - Computes SNR as per - https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 + Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 + for the given timesteps using the provided noise scheduler. + + Args: + noise_scheduler (`NoiseScheduler`): + An object containing the noise schedule parameters, specifically `alphas_cumprod`, which is used + to compute the SNR values. + timesteps (`torch.Tensor`): + A tensor of timesteps for which the SNR is computed. + + Returns: + `torch.Tensor`: A tensor containing the computed SNR values for each timestep. """ alphas_cumprod = noise_scheduler.alphas_cumprod sqrt_alphas_cumprod = alphas_cumprod**0.5 From 8eb523d1cb8bd2a4acf0b1f2430fda09f0fb8ad6 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 2 Nov 2024 07:31:18 +0530 Subject: [PATCH 2/2] quality --- src/diffusers/training_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py index 6bb8b11fa42e..2474ed5c2114 100644 --- a/src/diffusers/training_utils.py +++ b/src/diffusers/training_utils.py @@ -59,13 +59,14 @@ def set_seed(seed: int): def compute_snr(noise_scheduler, timesteps): """ - Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 + Computes SNR as per + https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 for the given timesteps using the provided noise scheduler. Args: noise_scheduler (`NoiseScheduler`): - An object containing the noise schedule parameters, specifically `alphas_cumprod`, which is used - to compute the SNR values. + An object containing the noise schedule parameters, specifically `alphas_cumprod`, which is used to compute + the SNR values. timesteps (`torch.Tensor`): A tensor of timesteps for which the SNR is computed.