flaxdiff 0.1.1__py3-none-any.whl

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,216 @@
+import orbax.checkpoint
+import tqdm
+from flax import linen as nn
+import jax
+from typing import Callable
+from dataclasses import field
+import jax.numpy as jnp
+from clu import metrics
+from flax.training import train_state  # Useful dataclass to keep train state
+import optax
+from flax import struct                # Flax dataclasses
+import time
+import os
+import orbax
+from flax.training import orbax_utils
+
+from ..schedulers import NoiseScheduler
+from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+
+@struct.dataclass
+class Metrics(metrics.Collection):
+  loss: metrics.Average.from_output('loss') # type: ignore
+
+class ModelState():
+    model: nn.Module
+    params: dict
+    noise_schedule: NoiseScheduler
+    model_output_transform: DiffusionPredictionTransform
+
+# Define the TrainState with EMA parameters
+class TrainState(train_state.TrainState):
+    rngs: jax.random.PRNGKey
+    ema_params: dict
+
+    def get_random_key(self):
+        rngs, subkey = jax.random.split(self.rngs)
+        return self.replace(rngs=rngs), subkey
+
+    def apply_ema(self, decay: float=0.999):
+        new_ema_params = jax.tree_util.tree_map(
+            lambda ema, param: decay * ema + (1 - decay) * param,
+            self.ema_params,
+            self.params,
+        )
+        return self.replace(ema_params=new_ema_params)
+
+class DiffusionTrainer:
+    state : TrainState
+    best_state : TrainState
+    best_loss : float
+    model : nn.Module
+    noise_schedule : NoiseScheduler
+    model_output_transform:DiffusionPredictionTransform
+    ema_decay:float = 0.999
+    
+    def __init__(self, 
+                 model:nn.Module, 
+                 optimizer: optax.GradientTransformation,
+                 noise_schedule:NoiseScheduler,
+                 rngs:jax.random.PRNGKey,
+                 train_state:TrainState=None,
+                 name:str="Diffusion",
+                 load_from_checkpoint:bool=False,
+                 param_transforms:Callable=None,
+                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform(),
+                 loss_fn=optax.l2_loss,
+                 ):
+        self.model = model
+        self.noise_schedule = noise_schedule
+        self.name = name
+        self.model_output_transform = model_output_transform
+        self.loss_fn = loss_fn
+
+        checkpointer = orbax.checkpoint.PyTreeCheckpointer()
+        options = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=4, create=True)
+        self.checkpointer = orbax.checkpoint.CheckpointManager(self.checkpoint_path(), checkpointer, options)
+
+        if load_from_checkpoint:
+            params = self.load()
+        else:
+            params = None
+
+        if train_state == None:
+            self.init_state(optimizer, rngs, params=params, model=model, param_transforms=param_transforms)
+        else:
+            self.state = train_state
+            self.best_state = train_state
+            self.best_loss = 1e9
+
+    def init_state(self, 
+                   optimizer: optax.GradientTransformation, 
+                   rngs:jax.random.PRNGKey,
+                   params:dict=None,
+                   model:nn.Module=None,
+                     param_transforms:Callable=None,
+                     batch_size=16,
+                    image_size=64
+                   ):
+        inp = jnp.ones((batch_size, image_size, image_size, 3))
+        temb = jnp.ones((batch_size,))
+        rngs, subkey = jax.random.split(rngs)
+        if params == None:
+            params = model.init(subkey, inp, temb)
+        if param_transforms is not None:
+            params = param_transforms(params)
+        self.best_loss = 1e9
+        self.state = TrainState.create(
+            apply_fn=model.apply,
+            params=params,
+            ema_params=params,
+            tx=optimizer,
+            rngs=rngs,
+        )
+        self.best_state = self.state
+
+    def checkpoint_path(self):
+        experiment_name = self.name
+        path = os.path.join(os.path.abspath('./models'), experiment_name)
+        if not os.path.exists(path):
+            os.makedirs(path)
+        return path
+
+    def load(self):
+        step = self.checkpointer.latest_step()
+        print("Loading model from checkpoint", step)
+        ckpt = self.checkpointer.restore(step)
+        state = ckpt['state']
+        # Convert the state to a TrainState
+        self.best_loss = ckpt['best_loss']
+        print(f"Loaded model from checkpoint at step {step}", ckpt['best_loss'])
+        return state.get('params', None)#, ckpt.get('model', None)
+
+    def save(self, epoch=0, best=False):
+        print(f"Saving model at epoch {epoch}")
+        state = self.best_state if best else self.state
+        # filename = os.path.join(self.checkpoint_path(), f'model_{epoch}' if not best else 'best_model')
+        ckpt = {
+            'model': self.model,
+            'state': state,
+            'best_loss': self.best_loss
+        }
+        save_args = orbax_utils.save_args_from_target(ckpt)
+        self.checkpointer.save(epoch, ckpt, save_kwargs={'save_args': save_args})
+
+    def summary(self, image_size=64):
+        inp = jnp.ones((1, image_size, image_size, 3))
+        temb = jnp.ones((1,))
+        print(self.model.tabulate(jax.random.key(0), inp, temb, console_kwargs={"width": 200, "force_jupyter":True, }))
+
+    def _define_train_step(self):
+        noise_schedule = self.noise_schedule
+        model = self.model
+        model_output_transform = self.model_output_transform
+        loss_fn = self.loss_fn
+        @jax.jit
+        def train_step(state:TrainState, batch):
+            """Train for a single step."""
+            images = batch
+            noise_level, state = noise_schedule.generate_timesteps(images.shape[0], state)
+            state, rngs = state.get_random_key()
+            noise:jax.Array = jax.random.normal(rngs, shape=images.shape)
+            rates = noise_schedule.get_rates(noise_level)
+            noisy_images, c_in, expected_output = model_output_transform.forward_diffusion(images, noise, rates)
+            def model_loss(params):
+                preds = model.apply(params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level))
+                preds = model_output_transform.pred_transform(noisy_images, preds, rates)
+                nloss = loss_fn(preds, expected_output)
+                # nloss = jnp.mean(nloss, axis=1)
+                nloss *= noise_schedule.get_weights(noise_level)
+                nloss = jnp.mean(nloss)
+                loss = nloss
+                return loss
+            loss, grads = jax.value_and_grad(model_loss)(state.params)
+            state = state.apply_gradients(grads=grads) 
+            state = state.apply_ema(self.ema_decay)
+            return state, loss
+        return train_step
+    
+    def _define_compute_metrics(self):
+        @jax.jit
+        def compute_metrics(state:TrainState, expected, pred):
+            loss = jnp.mean(jnp.square(pred - expected))
+            metric_updates = state.metrics.single_from_model_output(loss=loss)
+            metrics = state.metrics.merge(metric_updates)
+            state = state.replace(metrics=metrics)
+            return state
+        return compute_metrics
+
+    def fit(self, data, steps_per_epoch, epochs):
+        data = iter(data)
+        train_step = self._define_train_step()
+        compute_metrics = self._define_compute_metrics()
+        state = self.state
+        for epoch in range(epochs):
+            print(f"\nEpoch {epoch+1}/{epochs}")
+            start_time = time.time()
+            epoch_loss = 0
+            with tqdm.tqdm(total=steps_per_epoch, desc=f'\t\tEpoch {epoch+1}', ncols=100, unit='step') as pbar:
+                for i in range(steps_per_epoch):
+                    batch = next(data)
+                    state, loss = train_step(state, batch)
+                    epoch_loss += loss
+                    if i % 100 == 0:
+                        pbar.set_postfix(loss=f'{loss:.4f}')
+                        pbar.update(100)
+            end_time = time.time()
+            self.state = state
+            total_time = end_time - start_time
+            avg_time_per_step = total_time / steps_per_epoch
+            avg_loss = epoch_loss / steps_per_epoch
+            if avg_loss < self.best_loss:
+                self.best_loss = avg_loss
+                self.best_state = state
+                self.save(epoch, best=True)
+            print(f"\n\tEpoch {epoch+1} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}")
+        return self.state

flaxdiff/utils.py

@@ -0,0 +1,89 @@
+import jax
+import jax.numpy as jnp
+import flax.struct as struct
+import flax.linen as nn
+from typing import Any
+
+class MarkovState(struct.PyTreeNode):
+    pass
+
+class RandomMarkovState(MarkovState):
+    rng: jax.random.PRNGKey
+
+    def get_random_key(self):
+        rng, subkey = jax.random.split(self.rng)
+        return RandomMarkovState(rng), subkey
+
+def clip_images(images, clip_min=-1, clip_max=1):
+    return jnp.clip(images, clip_min, clip_max)
+
+class RMSNorm(nn.Module):
+    """
+    From "Root Mean Square Layer Normalization" by https://arxiv.org/abs/1910.07467
+
+    Adapted from flax.linen.LayerNorm
+    """
+
+    epsilon: float = 1e-6
+    dtype: Any = jnp.float32
+    param_dtype: Any = jnp.float32
+    use_scale: bool = True
+    scale_init: Any = jax.nn.initializers.ones
+
+    @nn.compact
+    def __call__(self, x):
+        reduction_axes = (-1,)
+        feature_axes = (-1,)
+
+        rms_sq = self._compute_rms_sq(x, reduction_axes)
+
+        return self._normalize(
+            self,
+            x,
+            rms_sq,
+            reduction_axes,
+            feature_axes,
+            self.dtype,
+            self.param_dtype,
+            self.epsilon,
+            self.use_scale,
+            self.scale_init,
+        )
+
+    def _compute_rms_sq(self, x, axes):
+        x = jnp.asarray(x, jnp.promote_types(jnp.float32, jnp.result_type(x)))
+        rms_sq = jnp.mean(jax.lax.square(x), axes)
+        return rms_sq
+
+    def _normalize(
+        self,
+        mdl,
+        x,
+        rms_sq,
+        reduction_axes,
+        feature_axes,
+        dtype,
+        param_dtype,
+        epsilon,
+        use_scale,
+        scale_init,
+    ):
+        reduction_axes = nn.normalization._canonicalize_axes(x.ndim, reduction_axes)
+        feature_axes = nn.normalization._canonicalize_axes(x.ndim, feature_axes)
+        stats_shape = list(x.shape)
+        for axis in reduction_axes:
+            stats_shape[axis] = 1
+        rms_sq = rms_sq.reshape(stats_shape)
+        feature_shape = [1] * x.ndim
+        reduced_feature_shape = []
+        for ax in feature_axes:
+            feature_shape[ax] = x.shape[ax]
+            reduced_feature_shape.append(x.shape[ax])
+        mul = jax.lax.rsqrt(rms_sq + epsilon)
+        if use_scale:
+            scale = mdl.param(
+                "scale", scale_init, reduced_feature_shape, param_dtype
+            ).reshape(feature_shape)
+            mul *= scale
+        y = mul * x
+        return jnp.asarray(y, dtype)