flaxdiff 0.1.1__py3-none-any.whl

flaxdiff/samplers/common.py

@@ -0,0 +1,113 @@
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+import tqdm
+from typing import Union
+from ..schedulers import NoiseScheduler
+from ..utils import RandomMarkovState, MarkovState, clip_images
+from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+
+class DiffusionSampler():
+    model:nn.Module
+    noise_schedule:NoiseScheduler
+    params:dict
+    model_output_transform:DiffusionPredictionTransform
+
+    def __init__(self, model:nn.Module, params:dict,  
+                 noise_schedule:NoiseScheduler, 
+                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform()):
+        self.model = model
+        self.noise_schedule = noise_schedule
+        self.params = params
+        self.model_output_transform = model_output_transform
+
+        @jax.jit
+        def sample_model(x_t, t):
+            rates = self.noise_schedule.get_rates(t)
+            c_in = self.model_output_transform.get_input_scale(rates)
+            model_output = self.model.apply(self.params, *self.noise_schedule.transform_inputs(x_t * c_in, t))
+            x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+            return x_0, eps, model_output
+        
+        self.sample_model = sample_model
+
+    # Used to sample from the diffusion model
+    def sample_step(self, current_samples:jnp.ndarray, current_step, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
+        # First clip the noisy images
+        # pred_images = clip_images(pred_images)
+        step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.int32)
+        current_step = step_ones * current_step
+        next_step = step_ones * next_step
+        pred_images, pred_noise, _ = self.sample_model(current_samples, current_step)
+        # plotImages(pred_images)
+        new_samples, state =  self.take_next_step(current_samples=current_samples, reconstructed_samples=pred_images, 
+                             pred_noise=pred_noise, current_step=current_step, next_step=next_step, state=state)
+        return new_samples, state
+
+    def take_next_step(self, current_samples, reconstructed_samples, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
+        # estimate the q(x_{t-1} | x_t, x_0). 
+        # pred_images is x_0, noisy_images is x_t, steps is t
+        return NotImplementedError
+    
+    def scale_steps(self, steps):
+        scale_factor = self.noise_schedule.max_timesteps / 1000
+        return steps * scale_factor
+
+    def get_steps(self, start_step, end_step, diffusion_steps):
+        step_range = start_step - end_step
+        if diffusion_steps is None or diffusion_steps == 0:
+            diffusion_steps = start_step - end_step
+        diffusion_steps = min(diffusion_steps, step_range)
+        steps = jnp.linspace(end_step, start_step, diffusion_steps, dtype=jnp.int16)[::-1]
+        return steps
+    
+    def get_initial_samples(self, num_images, rngs:jax.random.PRNGKey, start_step, image_size=64):
+        start_step = self.scale_steps(start_step)
+        alpha_n, sigma_n = self.noise_schedule.get_rates(start_step)
+        variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2) 
+        return jax.random.normal(rngs, (num_images, image_size, image_size, 3)) * variance
+
+    def generate_images(self,
+                        num_images=16, 
+                        diffusion_steps=1000, 
+                        start_step:int = None,
+                        end_step:int = 0,
+                        steps_override=None,
+                        priors=None, 
+                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42))) -> jnp.ndarray:
+        if priors is None:
+            rngstate, newrngs = rngstate.get_random_key()
+            samples = self.get_initial_samples(num_images, newrngs, start_step)
+        else:
+            print("Using priors")
+            samples = priors
+
+        @jax.jit
+        def sample_step(state:RandomMarkovState, samples, current_step, next_step):
+            samples, state = self.sample_step(current_samples=samples,
+                                              current_step=current_step, 
+                                              state=state, next_step=next_step)
+            return samples, state
+
+        if start_step is None:
+            start_step = self.noise_schedule.max_timesteps
+
+        if steps_override is not None:
+            steps = steps_override
+        else:
+            steps = self.get_steps(start_step, end_step, diffusion_steps)
+
+        # print("Sampling steps", steps)
+        for i in tqdm.tqdm(range(0, len(steps))):
+            current_step = self.scale_steps(steps[i])
+            next_step = self.scale_steps(steps[i+1] if i+1 < len(steps) else 0)
+            if i != len(steps) - 1:
+                # print("normal step")
+                samples, rngstate = sample_step(rngstate, samples, current_step, next_step)
+            else:
+                # print("last step")
+                step_ones = jnp.ones((num_images, ), dtype=jnp.int32)
+                samples, _, _ = self.sample_model(samples, current_step * step_ones)
+        samples = clip_images(samples)
+        return samples

flaxdiff/samplers/ddim.py

@@ -0,0 +1,10 @@
+import jax.numpy as jnp
+from .common import DiffusionSampler
+from ..utils import MarkovState
+
+class DDIMSampler(DiffusionSampler):
+    def take_next_step(self, 
+                 current_samples, reconstructed_samples, 
+                 pred_noise, current_step, state:MarkovState, next_step=None) -> tuple[jnp.ndarray, MarkovState]:
+        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,43 @@
+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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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)
+        betas = (1 - signal_ratio_squared)
+        gamma = jnp.sqrt(noise_ratio_squared * betas)
+        
+        next_samples = next_signal_rate * reconstructed_samples + pred_noise_coeff * pred_noise + noise * gamma
+        # pred_noise_coeff = ((next_noise_rate ** 2) * current_signal_rate) / (current_noise_rate * next_signal_rate)
+        # next_samples =  (2 - jnp.sqrt(1 - betas)) * current_samples - betas * (pred_noise / current_noise_rate) + noise * gamma#jnp.sqrt(betas)
+        # next_samples = (1 / (jnp.sqrt(1 - betas) + 1.e-24)) * (current_samples - betas * (pred_noise / current_noise_rate)) + noise * gamma
+        return next_samples, state

flaxdiff/samplers/euler.py

@@ -0,0 +1,59 @@
+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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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,28 @@
+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, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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)
+        
+        # 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,60 @@
+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 _renoise(self, 
+                 current_samples, reconstructed_samples, 
+                 pred_noise, current_step, state:RandomMarkovState, next_step=None) -> 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
+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) -> tuple[jnp.ndarray, RandomMarkovState]:
+            t = self.noise_schedule.get_timesteps(sigma)
+            x_0, eps, _ = self.sample_model(x_t, t)
+            return eps, state
+        
+        self.get_derivative = get_derivative
+
+    def sample_step(self, current_samples:jnp.ndarray, current_step, next_step, state:RandomMarkovState=None) -> tuple[jnp.ndarray, RandomMarkovState]:
+        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)
+        k2, state = self.get_derivative(current_samples + 0.5 * k1 * dt, current_sigma + 0.5 * dt, state)
+        k3, state = self.get_derivative(current_samples + 0.5 * k2 * dt, current_sigma + 0.5 * dt, state)
+        k4, state = self.get_derivative(current_samples + k3 * dt, current_sigma + dt, state)
+
+        next_samples = current_samples + ((k1 + 2 * k2 + 2 * k3 + k4) / 6) * dt
+        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)