flaxdiff 0.1.36.1__py3-none-any.whl → 0.1.36.3__py3-none-any.whl

flaxdiff/samplers/ddim.py

@@ -0,0 +1,10 @@
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import MarkovState, RandomMarkovState
+
+class DDIMSampler(DiffusionSampler):
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step)
+        return reconstructed_samples * next_signal_rate + pred_noise * next_noise_rate, state
+    

flaxdiff/samplers/ddpm.py

@@ -0,0 +1,37 @@
+import jax
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import MarkovState, RandomMarkovState
+class DDPMSampler(DiffusionSampler):
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        mean = self.noise_schedule.get_posterior_mean(reconstructed_samples, current_samples, current_step)
+        variance = self.noise_schedule.get_posterior_variance(steps=current_step)
+        
+        state, rng = state.get_random_key()
+        # Now sample from the posterior
+        noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32)
+
+        return mean + noise * variance, state
+    
+    def generate_images(self, num_images=16, diffusion_steps=1000, start_step: int = None, *args, **kwargs):
+        return super().generate_images(num_images=num_images, diffusion_steps=diffusion_steps, start_step=start_step, *args, **kwargs)
+    
+class SimpleDDPMSampler(DiffusionSampler):
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        state, rng = state.get_random_key()
+        noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32)
+
+        # Compute noise rates and signal rates only once
+        current_signal_rate, current_noise_rate = self.noise_schedule.get_rates(current_step)
+        next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step)
+        
+        pred_noise_coeff = ((next_noise_rate ** 2) * current_signal_rate) / (current_noise_rate * next_signal_rate)
+        
+        noise_ratio_squared = (next_noise_rate ** 2) / (current_noise_rate ** 2)
+        signal_ratio_squared = (current_signal_rate ** 2) / (next_signal_rate ** 2)
+        gamma = jnp.sqrt(noise_ratio_squared * (1 - signal_ratio_squared))
+        
+        next_samples = next_signal_rate * reconstructed_samples + pred_noise_coeff * pred_noise + noise * gamma
+        return next_samples, state

flaxdiff/samplers/euler.py

@@ -0,0 +1,56 @@
+import jax
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import RandomMarkovState
+
+class EulerSampler(DiffusionSampler):
+    # Basically a DDIM Sampler but parameterized as an ODE
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        dt = next_sigma - current_sigma
+        
+        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / (dt)
+        dx = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
+        next_samples = current_samples + dx * dt
+        return next_samples, state
+
+class SimplifiedEulerSampler(DiffusionSampler):
+    """
+    This is for networks with forward diffusion of the form x_{t+1} = x_t + sigma_t * epsilon_t
+    """
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        _, current_sigma = self.noise_schedule.get_rates(current_step)
+        _, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        dt = next_sigma - current_sigma
+        
+        dx = (current_samples - reconstructed_samples) / current_sigma
+        next_samples = current_samples + dx * dt
+        return next_samples, state
+    
+class EulerAncestralSampler(DiffusionSampler):
+    """
+    Similar to EulerSampler but with ancestral sampling
+    """
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        sigma_up = (next_sigma**2 * (current_sigma**2 - next_sigma**2) / current_sigma**2) ** 0.5
+        sigma_down = (next_sigma**2 - sigma_up**2) ** 0.5
+        
+        dt = sigma_down - current_sigma
+        
+        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / (next_sigma - current_sigma)
+        dx = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
+        
+        state, subkey = state.get_random_key()
+        dW = jax.random.normal(subkey, current_samples.shape) * sigma_up
+        
+        next_samples = current_samples + dx * dt + dW
+        return next_samples, state

flaxdiff/samplers/heun_sampler.py

@@ -0,0 +1,27 @@
+import jax
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import RandomMarkovState
+
+class HeunSampler(DiffusionSampler):
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        # Get the noise and signal rates for the current and next steps
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        dt = next_sigma - current_sigma
+        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt
+
+        dx_0 = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
+        next_samples_0 = current_samples + dx_0 * dt
+        
+        # Recompute x_0 and eps at the first estimate to refine the derivative
+        estimated_x_0, _, _ = self.sample_model(next_samples_0, next_step, *model_conditioning_inputs)
+        
+        # Estimate the refined derivative using the midpoint (Heun's method)
+        dx_1 = (next_samples_0 - x_0_coeff * estimated_x_0) / next_sigma
+        # Compute the final next samples by averaging the initial and refined derivatives
+        final_next_samples = current_samples + 0.5 * (dx_0 + dx_1) * dt
+        
+        return final_next_samples, state

flaxdiff/samplers/multistep_dpm.py

@@ -0,0 +1,59 @@
+import jax
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import RandomMarkovState
+
+class MultiStepDPM(DiffusionSampler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.history = []
+
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        # Get the noise and signal rates for the current and next steps
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        dt = next_sigma - current_sigma
+
+        def first_order(current_noise, current_sigma):
+            dx = current_noise
+            return dx
+        
+        def second_order(current_noise, current_sigma, last_noise, last_sigma):
+            dx_2 = (current_noise - last_noise) / (current_sigma - last_sigma)
+            return dx_2
+        
+        def third_order(current_noise, current_sigma, last_noise, last_sigma, second_last_noise, second_last_sigma):
+            dx_2 = second_order(current_noise, current_sigma, last_noise, last_sigma)
+            dx_2_last = second_order(last_noise, last_sigma, second_last_noise, second_last_sigma)
+
+            dx_3 = (dx_2 - dx_2_last) / (0.5 * ((current_sigma + last_sigma) - (last_sigma + second_last_sigma)))
+            
+            return dx_3
+
+        if len(self.history) == 0:
+            # First order only
+            dx_1 = first_order(pred_noise, current_sigma)
+            next_samples = current_samples + dx_1 * dt
+        elif len(self.history) == 1:
+            # First + Second order
+            dx_1 = first_order(pred_noise, current_sigma)
+            last_step = self.history[-1]
+            dx_2 = second_order(pred_noise, current_sigma, last_step['eps'], last_step['sigma'])
+            next_samples = current_samples + dx_1 * dt + 0.5 * dx_2 * dt**2
+        else:
+            # First + Second + Third order
+            last_step = self.history[-1]
+            second_last_step = self.history[-2]
+
+            dx_1 = first_order(pred_noise, current_sigma)
+            dx_2 = second_order(pred_noise, current_sigma, last_step['eps'], last_step['sigma'])
+            dx_3 = third_order(pred_noise, current_sigma, last_step['eps'], last_step['sigma'], second_last_step['eps'], second_last_step['sigma'])
+            next_samples = current_samples + (dx_1 * dt) + (0.5 * dx_2 * dt**2) + ((1/6) * dx_3 * dt**3)
+
+        self.history.append({
+            "eps": pred_noise,
+            "sigma" : current_sigma,
+        })
+        return next_samples, state

flaxdiff/samplers/rk4_sampler.py

@@ -0,0 +1,34 @@
+import jax
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import RandomMarkovState, MarkovState
+from ..schedulers import GeneralizedNoiseScheduler
+
+class RK4Sampler(DiffusionSampler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert issubclass(type(self.noise_schedule), GeneralizedNoiseScheduler), "Noise schedule must be a GeneralizedNoiseScheduler"
+        @jax.jit
+        def get_derivative(x_t, sigma, state:RandomMarkovState, model_conditioning_inputs) -> tuple[jnp.ndarray, RandomMarkovState]:
+            t = self.noise_schedule.get_timesteps(sigma)
+            x_0, eps, _ = self.sample_model(x_t, t, *model_conditioning_inputs)
+            return eps, state
+        
+        self.get_derivative = get_derivative
+
+    def sample_step(self, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
+        step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.int32)
+        current_step = step_ones * current_step
+        next_step = step_ones * next_step
+        _, current_sigma = self.noise_schedule.get_rates(current_step)
+        _, next_sigma = self.noise_schedule.get_rates(next_step)
+
+        dt = next_sigma - current_sigma
+
+        k1, state = self.get_derivative(current_samples, current_sigma, state, model_conditioning_inputs)
+        k2, state = self.get_derivative(current_samples + 0.5 * k1 * dt, current_sigma + 0.5 * dt, state, model_conditioning_inputs)
+        k3, state = self.get_derivative(current_samples + 0.5 * k2 * dt, current_sigma + 0.5 * dt, state, model_conditioning_inputs)
+        k4, state = self.get_derivative(current_samples + k3 * dt, current_sigma + dt, state, model_conditioning_inputs)
+
+        next_samples = current_samples + (((k1 + 2 * k2 + 2 * k3 + k4) * dt) / 6)
+        return next_samples, state

flaxdiff/schedulers/__init__.py

@@ -0,0 +1,6 @@
+from .discrete import DiscreteNoiseScheduler
+from .common import NoiseScheduler, GeneralizedNoiseScheduler
+from .cosine import CosineNoiseSchedule, ContinuousNoiseScheduler, CosineGeneralNoiseScheduler
+from .linear import LinearNoiseSchedule
+from .sqrt import SqrtContinuousNoiseScheduler
+from .karras import KarrasVENoiseScheduler, SimpleExpNoiseScheduler, EDMNoiseScheduler

flaxdiff/schedulers/common.py

@@ -0,0 +1,98 @@
+import jax
+import jax.numpy as jnp
+from typing import Union
+from ..utils import RandomMarkovState  
+
+class NoiseScheduler():
+    def __init__(self, timesteps,
+                    dtype=jnp.float32,
+                    clip_min=-1.0,
+                    clip_max=1.0,
+                    *args, **kwargs):
+        self.max_timesteps = timesteps
+        self.dtype = dtype
+        self.clip_min = clip_min
+        self.clip_max = clip_max
+        if type(timesteps) == int and timesteps > 1:
+            timestep_generator = lambda rng, batch_size, max_timesteps = timesteps: jax.random.randint(rng, (batch_size,), 0, max_timesteps)
+        else:
+            timestep_generator = lambda rng, batch_size, max_timesteps = timesteps: jax.random.uniform(rng, (batch_size,), minval=0, maxval=max_timesteps)
+        self.timestep_generator = timestep_generator
+
+    def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
+        state, rng = state.get_random_key()
+        timesteps = self.timestep_generator(rng, batch_size, self.max_timesteps)
+        return timesteps, state
+    
+    def get_weights(self, steps):
+        raise NotImplementedError
+    
+    def reshape_rates(self, rates:tuple[jnp.ndarray, jnp.ndarray], shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        signal_rates, noise_rates = rates
+        signal_rates = jnp.reshape(signal_rates, shape)
+        noise_rates = jnp.reshape(noise_rates, shape)
+        return signal_rates, noise_rates
+    
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        raise NotImplementedError
+    
+    def add_noise(self, images, noise, steps) -> jnp.ndarray:
+        signal_rates, noise_rates = self.get_rates(steps)
+        return signal_rates * images + noise_rates * noise
+    
+    def remove_all_noise(self, noisy_images, noise, steps, clip_denoised=True, rates=None):
+        signal_rates, noise_rates = self.get_rates(steps)
+        x_0 = (noisy_images - noise * noise_rates) / signal_rates
+        return x_0
+    
+    def transform_inputs(self, x, steps):
+        return x, steps
+    
+    def get_posterior_mean(self, x_0, x_t, steps):
+        raise NotImplementedError
+    
+    def get_posterior_variance(self, steps, shape=(-1, 1, 1, 1)):
+        raise NotImplementedError
+
+    def get_max_variance(self):
+        alpha_n, sigma_n = self.get_rates(self.max_timesteps)
+        variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2) 
+        return variance
+
+class GeneralizedNoiseScheduler(NoiseScheduler):
+    """
+    As per the generalization presented in the paper
+    "Elucidating the Design Space of Diffusion-Based
+    Generative Models" by Tero Karras et al.
+    Basically the signal rate shall always be 1, and the model
+    input itself shall be scaled to match the noise rate
+    """
+    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80.0, sigma_data=1, *args, **kwargs):
+        super().__init__(timesteps, *args, **kwargs)
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.sigma_data = sigma_data
+    
+    def get_weights(self, steps, shape=(-1, 1, 1, 1)):
+        sigma = self.get_sigmas(steps)
+        return (1 + (1 / (1 + ((1 - sigma ** 2)/(sigma ** 2)))) / (self.sigma_max ** 2)).reshape(shape)
+    
+    def get_sigmas(self, steps) -> jnp.ndarray:
+        raise NotImplementedError("This method should be implemented in the subclass")
+    
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        sigmas = self.get_sigmas(steps)
+        signal_rates = 1
+        noise_rates = sigmas
+        return self.reshape_rates((signal_rates, noise_rates), shape=shape)
+    
+    def transform_inputs(self, x, steps, num_discrete_chunks=1000):
+        sigmas_discrete = (steps / self.max_timesteps) * num_discrete_chunks
+        sigmas_discrete = sigmas_discrete.astype(jnp.int32)
+        return x, sigmas_discrete
+    
+    def get_timesteps(self, sigmas):
+        """
+        Inverse of the get_sigmas method
+        """
+        raise NotImplementedError("This method should be implemented in the subclass")

flaxdiff/schedulers/continuous.py

@@ -0,0 +1,12 @@
+import jax
+import jax.numpy as jnp
+from typing import Union
+from ..utils import RandomMarkovState  
+from .common import NoiseScheduler
+
+class ContinuousNoiseScheduler(NoiseScheduler):
+    """
+    General Continuous Noise Scheduler
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(timesteps=1, *args, **kwargs)

flaxdiff/schedulers/cosine.py

@@ -0,0 +1,40 @@
+import math
+import numpy as np
+import jax.numpy as jnp
+from .discrete import DiscreteNoiseScheduler
+from .continuous import ContinuousNoiseScheduler
+from .common import GeneralizedNoiseScheduler
+
+def cosine_beta_schedule(timesteps, start_angle=0.008, end_angle=0.999):
+    ts = np.linspace(0, 1, timesteps + 1, dtype=np.float64)
+    alphas_bar = np.cos((ts + start_angle) / (1 + start_angle) * np.pi /2) ** 2
+    alphas_bar = alphas_bar/alphas_bar[0]
+    betas = 1 - (alphas_bar[1:] / alphas_bar[:-1])
+    return np.clip(betas, 0, end_angle)
+
+class CosineNoiseSchedule(DiscreteNoiseScheduler):
+    def __init__(self, timesteps, beta_start=0.008, beta_end=0.999, *args, **kwargs):
+        super().__init__(timesteps, beta_start, beta_end, schedule_fn=cosine_beta_schedule, *args, **kwargs)
+
+class CosineGeneralNoiseScheduler(GeneralizedNoiseScheduler):
+    def __init__(self, sigma_min=0.02, sigma_max=80.0, kappa=1.0, *args, **kwargs):
+        super().__init__(timesteps=1, sigma_min=sigma_min, sigma_max=sigma_max, *args, **kwargs)
+        self.kappa = kappa
+        logsnr_max = 2 * (math.log(self.kappa) - math.log(self.sigma_max))
+        self.theta_max = math.atan(math.exp(-0.5 * logsnr_max))
+        logsnr_min = 2 * (math.log(self.kappa) - math.log(self.sigma_min))
+        self.theta_min = math.atan(math.exp(-0.5 * logsnr_min))
+    
+    def get_sigmas(self, steps):
+        return jnp.tan(self.theta_min + steps * (self.theta_max - self.theta_min)) / self.kappa
+    
+class CosineContinuousNoiseScheduler(ContinuousNoiseScheduler):
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        signal_rates = jnp.cos((jnp.pi * steps) / (2 * self.max_timesteps))
+        noise_rates = jnp.sin((jnp.pi * steps) / (2 * self.max_timesteps))
+        return self.reshape_rates((signal_rates, noise_rates), shape=shape)
+    
+    def get_weights(self, steps):
+        alpha, sigma = self.get_rates(steps, shape=())
+        return 1 / (1 + (alpha ** 2 / sigma ** 2))
+    

flaxdiff/schedulers/discrete.py

@@ -0,0 +1,74 @@
+import jax
+import jax.numpy as jnp
+from typing import Union
+from ..utils import RandomMarkovState  
+from .common import NoiseScheduler
+
+class DiscreteNoiseScheduler(NoiseScheduler):
+    """
+    Variance Preserving Noise Scheduler
+    signal_rate**2 + noise_rate**2 = 1
+    """
+    def __init__(self, timesteps,
+                    beta_start=0.0001,
+                    beta_end=0.02,
+                    schedule_fn=None, 
+                    p2_loss_weight_k:float=1,
+                    p2_loss_weight_gamma:float=1,
+                    *args, **kwargs):
+        super().__init__(timesteps, *args, **kwargs)
+        betas = schedule_fn(timesteps, beta_start, beta_end)
+        alphas = 1 - betas
+        alpha_cumprod = jnp.cumprod(alphas, axis=0)
+        alpha_cumprod_prev = jnp.append(1.0, alpha_cumprod[:-1])
+        
+        self.betas = jnp.array(betas, dtype=jnp.float32)
+        self.alphas = alphas.astype(jnp.float32)
+        self.alpha_cumprod = alpha_cumprod.astype(jnp.float32)
+        self.alpha_cumprod_prev = alpha_cumprod_prev.astype(jnp.float32)
+
+        self.sqrt_alpha_cumprod = jnp.sqrt(alpha_cumprod).astype(jnp.float32)
+        self.sqrt_one_minus_alpha_cumprod = jnp.sqrt(1 - alpha_cumprod).astype(jnp.float32)
+
+        posterior_variance = (betas * (1 - alpha_cumprod_prev) / (1 - alpha_cumprod))
+        self.posterior_variance = posterior_variance.astype(jnp.float32)
+        self.posterior_log_variance_clipped = (jnp.log(jnp.maximum(posterior_variance, 1e-20))).astype(jnp.float32)
+        
+        self.posterior_mean_coef1 = (betas * jnp.sqrt(alpha_cumprod_prev) / (1 - alpha_cumprod)).astype(jnp.float32)
+        self.posterior_mean_coef2 = ((1 - alpha_cumprod_prev) * jnp.sqrt(alphas) / (1 - alpha_cumprod)).astype(jnp.float32)
+
+        self.p2_loss_weights = self.get_p2_weights(p2_loss_weight_k, p2_loss_weight_gamma)
+    
+    def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
+        state, rng = state.get_random_key()
+        timesteps = jax.random.randint(rng, (batch_size,), 0, self.max_timesteps)
+        return timesteps, state
+    
+    def get_p2_weights(self, k, gamma):
+        return (k + self.alpha_cumprod / (1 - self.alpha_cumprod)) ** -gamma
+    
+    def get_weights(self, steps, shape=(-1, 1, 1, 1)):
+        steps = jnp.int16(steps)
+        return self.p2_loss_weights[steps].reshape(shape)
+
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        steps = jnp.int16(steps)
+        signal_rate = self.sqrt_alpha_cumprod[steps]
+        noise_rate = self.sqrt_one_minus_alpha_cumprod[steps]
+        signal_rate = jnp.reshape(signal_rate, shape)
+        noise_rate = jnp.reshape(noise_rate, shape)
+        return signal_rate, noise_rate
+    
+    def get_posterior_mean(self, x_0, x_t, steps):
+        steps = jnp.int16(steps)
+        x_0_coeff = self.posterior_mean_coef1[steps]
+        x_t_coeff = self.posterior_mean_coef2[steps]
+        x_0_coeff, x_t_coeff = self.reshape_rates((x_0_coeff, x_t_coeff))
+        mean = x_0_coeff * x_0 + x_t_coeff * x_t
+        return mean
+    
+    def get_posterior_variance(self, steps, shape=(-1, 1, 1, 1)):
+        steps = int(steps)
+        return jnp.exp(0.5 * self.posterior_log_variance_clipped[steps]).reshape(shape)
+
+        

flaxdiff/schedulers/exp.py

@@ -0,0 +1,13 @@
+import numpy as np
+from .discrete import DiscreteNoiseScheduler
+
+def exp_beta_schedule(timesteps, start_angle=0.008, end_angle=0.999):
+    ts = np.linspace(0, 1, timesteps + 1, dtype=np.float64)
+    alphas_bar = np.exp(ts * -12.0)
+    alphas_bar = alphas_bar/alphas_bar[0]
+    betas = 1 - (alphas_bar[1:] / alphas_bar[:-1])
+    return np.clip(betas, 0, end_angle)
+
+class ExpNoiseSchedule(DiscreteNoiseScheduler):
+    def __init__(self, timesteps, beta_start=0.008, beta_end=0.999, *args, **kwargs):
+        super().__init__(timesteps, beta_start, beta_end, schedule_fn=exp_beta_schedule, *args, **kwargs)

flaxdiff/schedulers/karras.py

@@ -0,0 +1,69 @@
+import jax.numpy as jnp
+from .common import GeneralizedNoiseScheduler
+import math
+import jax
+from ..utils import RandomMarkovState
+
+class KarrasVENoiseScheduler(GeneralizedNoiseScheduler):
+    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs):
+        super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs)
+        self.min_inv_rho = sigma_min ** (1 / rho)
+        self.max_inv_rho = sigma_max ** (1 / rho)
+        self.rho = rho
+
+    def get_sigmas(self, steps) -> jnp.ndarray:
+        # steps = jnp.int16(steps)
+        # return self.sigmas[steps]
+        ramp = 1 - steps / self.max_timesteps
+        sigmas = (self.max_inv_rho + ramp * (self.min_inv_rho - self.max_inv_rho)) ** self.rho
+        return sigmas
+
+    def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
+        sigma = self.get_sigmas(steps)
+        weights = ((sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2)
+        return weights.reshape(shape)
+    
+    def transform_inputs(self, x, steps, num_discrete_chunks=1000) -> tuple[jnp.ndarray, jnp.ndarray]:
+        sigmas = self.get_sigmas(steps)
+        # sigmas = (sigmas / self.sigma_max) * num_discrete_chunks
+        sigmas = jnp.log(sigmas) / 4
+        return x, sigmas
+    
+    def get_timesteps(self, sigmas:jnp.ndarray) -> jnp.ndarray:
+        sigmas = sigmas.reshape(-1)
+        inv_rho = sigmas ** (1 / self.rho)
+        ramp = ((inv_rho - self.max_inv_rho) / (self.min_inv_rho - self.max_inv_rho))
+        steps = 1 - ramp * self.max_timesteps
+        return steps
+    
+    def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
+        timesteps, state = super().generate_timesteps(batch_size, state)
+        timesteps = timesteps.astype(jnp.float32)
+        return timesteps, state
+    
+class SimpleExpNoiseScheduler(KarrasVENoiseScheduler):
+    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs):
+        super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs)
+        if type(timesteps) == int and timesteps > 1:
+            n = timesteps
+        else:
+            n = 1000
+        self.sigmas = jnp.exp(jnp.linspace(math.log(sigma_min), math.log(sigma_max), n))
+
+    def get_sigmas(self, steps) -> jnp.ndarray:
+        steps = jnp.int16(steps)
+        return self.sigmas[steps]
+
+class EDMNoiseScheduler(KarrasVENoiseScheduler):
+    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs):
+        super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs)
+
+    def get_sigmas(self, steps, std=1.2, mean=-1.2) -> jnp.ndarray:
+        space = steps / self.max_timesteps
+        # space = jax.scipy.special.erfinv(self.erf_sigma_min + steps * (self.erf_sigma_max - self.erf_sigma_min))
+        return jnp.exp(space * std + mean)
+    
+    def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
+        state, rng = state.get_random_key()
+        timesteps = jax.random.normal(rng, (batch_size,), dtype=jnp.float32)
+        return timesteps, state

flaxdiff/schedulers/linear.py

@@ -0,0 +1,14 @@
+import numpy as np
+from .discrete import DiscreteNoiseScheduler
+
+def linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02):
+    scale = 1000 / timesteps
+    beta_start = scale * beta_start
+    beta_end = scale * beta_end
+    betas = np.linspace(
+        beta_start, beta_end, timesteps, dtype=np.float64)
+    return betas
+
+class LinearNoiseSchedule(DiscreteNoiseScheduler):
+    def __init__(self, timesteps, beta_start=0.0001, beta_end=0.02, *args, **kwargs):
+        super().__init__(timesteps, beta_start, beta_end, schedule_fn=linear_beta_schedule, *args, **kwargs)

flaxdiff/schedulers/sqrt.py

@@ -0,0 +1,10 @@
+import numpy as np
+import jax.numpy as jnp
+from .discrete import DiscreteNoiseScheduler
+from .continuous import ContinuousNoiseScheduler
+
+class SqrtContinuousNoiseScheduler(ContinuousNoiseScheduler):
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        signal_rates = jnp.sqrt(1 - steps)
+        noise_rates = jnp.sqrt(steps)
+        return self.reshape_rates((signal_rates, noise_rates), shape=shape)

flaxdiff/trainer/__init__.py

@@ -0,0 +1,2 @@
+from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+from .diffusion_trainer import DiffusionTrainer, TrainState