added more scaling and same seeds for student and teacher

consistency_models/consistency_models.py

@@ -6,7 +6,7 @@
 from torch import Tensor, nn
 from tqdm.auto import tqdm
 
-from .utils import pad_dims_like
+from .utils import pad_dims_like, get_seeds, set_seeds
 
 
 def timesteps_schedule(
@@ -263,6 +263,42 @@ def output_scaling(
     """
     return (sigma_data * (sigma - sigma_min)) / (sigma_data**2 + sigma**2) ** 0.5
 
+def input_scaling(
+    sigma: Tensor, sigma_data: float = 0.5
+) -> Tensor:
+    """Computes the scaling value for the model's input.
+
+    Parameters
+    ----------
+    sigma : Tensor
+        Current standard deviation of the noise.
+    sigma_data : float, default=0.5
+        Standard deviation of the data.
+
+    Returns
+    -------
+    Tensor
+        Scaling value for the model's input.
+    """
+    return 1/(sigma**2+sigma_data**2)**0.5
+
+def sigma_scaling(
+    sigma: Tensor
+) -> Tensor:
+    """Computes the scaling value for the sigma.
+
+    Parameters
+    ----------
+    sigma : Tensor
+        Current standard deviation of the noise.
+
+    Returns
+    -------
+    Tensor
+        Scaling value for the sigma.
+    """
+    return sigma.log()/4
+
 
 def model_forward_wrapper(
     model: nn.Module,
@@ -297,12 +333,16 @@ def model_forward_wrapper(
     """
     c_skip = skip_scaling(sigma, sigma_data, sigma_min)
     c_out = output_scaling(sigma, sigma_data, sigma_min)
+    c_in = input_scaling(sigma, sigma_data)
+    sigma_scaled = input_scaling(sigma)
 
     # Pad dimensions as broadcasting will not work
     c_skip = pad_dims_like(c_skip, x)
     c_out = pad_dims_like(c_out, x)
+    c_in = pad_dims_like(c_in, x)
+    sigma_scaled = pad_dims_like(sigma_scaled, x)
 
-    return c_skip * x + c_out * model(x, sigma, **kwargs)
+    return c_skip * x + c_out * model(c_in * x, sigma_scaled, **kwargs)
 
 
 @dataclass
@@ -413,6 +453,7 @@ def __call__(
         next_sigmas = sigmas[timesteps + 1]
 
         next_noisy_x = x + pad_dims_like(next_sigmas, x) * noise
+        cpu_state, gpu_state = get_seeds(x.device)
         next_x = model_forward_wrapper(
             student_model,
             next_noisy_x,
@@ -424,6 +465,7 @@ def __call__(
 
         with torch.no_grad():
             current_noisy_x = x + pad_dims_like(current_sigmas, x) * noise
+            set_seeds(cpu_state, gpu_state, x.device)
             current_x = model_forward_wrapper(
                 teacher_model,
                 current_noisy_x,
@@ -529,6 +571,7 @@ def __call__(
         next_sigmas = sigmas[timesteps + 1]
 
         next_noisy_x = x + pad_dims_like(next_sigmas, x) * noise
+        cpu_state, gpu_state = get_seeds(x.device)
         next_x = model_forward_wrapper(
             model,
             next_noisy_x,
@@ -540,6 +583,7 @@ def __call__(
 
         with torch.no_grad():
             current_noisy_x = x + pad_dims_like(current_sigmas, x) * noise
+            set_seeds(cpu_state, gpu_state, x.device)
             current_x = model_forward_wrapper(
                 model,
                 current_noisy_x,

consistency_models/utils.py

@@ -1,5 +1,6 @@
 from typing import Iterator
 
+import torch
 from torch import Tensor, nn
 
 
@@ -61,3 +62,13 @@ def update_ema_model_(
     _update_ema_weights(ema_model.buffers(), online_model.buffers(), ema_decay_rate)
 
     return ema_model
+
+def get_seeds(device):
+    rng_state_cpu = torch.get_rng_state()
+    rng_state_gpu = torch.cuda.get_rng_state() if device != torch.device('cpu') else None
+    return rng_state_cpu, rng_state_gpu
+
+def set_seeds(rng_state_cpu, rng_state_gpu, device):
+    torch.set_rng_state(rng_state_cpu)
+    if device != torch.device('cpu'):
+        torch.cuda.set_rng_state(rng_state_gpu)