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

flaxdiff/samplers/ddpm.py

@@ -1,43 +0,0 @@
-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

@@ -1,59 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,60 +0,0 @@
-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

@@ -1,34 +0,0 @@
-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

@@ -1,6 +0,0 @@
-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

@@ -1,98 +0,0 @@
-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

@@ -1,12 +0,0 @@
-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

@@ -1,40 +0,0 @@
-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

@@ -1,74 +0,0 @@
-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

@@ -1,13 +0,0 @@
-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

@@ -1,69 +0,0 @@
-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

@@ -1,14 +0,0 @@
-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

@@ -1,10 +0,0 @@
-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

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