feat: updated samplers, added validation stage

Author: AshishKumar4

flaxdiff/data/sources/tfds.py

@@ -3,11 +3,20 @@
 import grain.python as pygrain
 from flaxdiff.utils import AutoTextTokenizer
 from typing import Dict
+import random
 
 # -----------------------------------------------------------------------------------------------#
 # Oxford flowers and other TFDS datasources -----------------------------------------------------#
 # -----------------------------------------------------------------------------------------------#
 
+PROMPT_TEMPLATES = [
+    "a photo of a {}",
+    "a photo of a {} flower",
+    "This is a photo of a {}",
+    "This is a photo of a {} flower",
+    "A photo of a {} flower",
+]
+
 def data_source_tfds(name, use_tf=True, split="all"):
     import tensorflow_datasets as tfds
     if use_tf:
@@ -23,7 +32,13 @@ def labelizer_oxford_flowers102(path):
         textlabels = [i.strip() for i in f.readlines()]
 
     def load_labels(sample):
-        return textlabels[int(sample['label'])]
+        raw = textlabels[int(sample['label'])]
+        # randomly select a prompt template
+        template = random.choice(PROMPT_TEMPLATES)
+        # format the template with the label
+        caption = template.format(raw)
+        # return the caption
+        return caption
     return load_labels
 
 def tfds_augmenters(image_scale, method):

flaxdiff/samplers/common.py

@@ -15,36 +15,76 @@ class DiffusionSampler():
 
     def __init__(self, model:nn.Module, params:dict,  
                  noise_schedule:NoiseScheduler, 
-                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform()):
+                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform(),
+                 guidance_scale:float = 0.0,
+                 null_labels_seq:jax.Array=None,
+                 autoencoder=None,
+                 image_size=256,
+                 autoenc_scale_reduction=8,
+                 autoenc_latent_channels=4,
+                 ):
         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.guidance_scale = guidance_scale
+        self.image_size = image_size
+        self.autoenc_scale_reduction = autoenc_scale_reduction
+        self.autoencoder = autoencoder
+        self.autoenc_latent_channels = autoenc_latent_channels
 
+        if self.guidance_scale > 0:
+            # Classifier free guidance
+            assert null_labels_seq is not None, "Null labels sequence is required for classifier-free guidance"
+            print("Using classifier-free guidance")
+            def sample_model(x_t, t, *additional_inputs):
+                # Concatenate unconditional and conditional inputs
+                x_t_cat = jnp.concatenate([x_t] * 2, axis=0)
+                t_cat = jnp.concatenate([t] * 2, axis=0)
+                rates_cat = self.noise_schedule.get_rates(t_cat)
+                c_in_cat = self.model_output_transform.get_input_scale(rates_cat)
+                
+                text_labels_seq, = additional_inputs
+                text_labels_seq = jnp.concatenate([text_labels_seq, jnp.broadcast_to(null_labels_seq, text_labels_seq.shape)], axis=0)
+                model_output = self.model.apply(self.params, *self.noise_schedule.transform_inputs(x_t_cat * c_in_cat, t_cat), text_labels_seq)
+                # Split model output into unconditional and conditional parts
+                model_output_cond, model_output_uncond = jnp.split(model_output, 2, axis=0)
+                model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
+                
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+        else:
+            # Unconditional sampling
+            def sample_model(x_t, t, *additional_inputs):
+                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), *additional_inputs)
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+            
+        # if jax.device_count() > 1:
+        #     mesh = jax.sharding.Mesh(jax.devices(), 'data')
+        #     sample_model = shard_map(sample_model, mesh=mesh, in_specs=(P('data'), P('data'), P('data')),
+        #                              out_specs=(P('data'), P('data'), P('data')))
+        sample_model = jax.jit(sample_model)
         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]:
+    def sample_step(self, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, 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)
+        pred_images, pred_noise, _ = self.sample_model(current_samples, current_step, *model_conditioning_inputs)
         # plotImages(pred_images)
+        # pred_images = clip_images(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)
+                             pred_noise=pred_noise, current_step=current_step, next_step=next_step, state=state,
+                             model_conditioning_inputs=model_conditioning_inputs
+                             )
         return new_samples, state
 
-    def take_next_step(self, current_samples, reconstructed_samples, 
+    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]:
         # estimate the q(x_{t-1} | x_t, x_0). 
         # pred_images is x_0, noisy_images is x_t, steps is t
@@ -62,11 +102,16 @@ def get_steps(self, start_step, end_step, diffusion_steps):
         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):
+    def get_initial_samples(self, num_images, rngs:jax.random.PRNGKey, start_step):
         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
+        image_size = self.image_size
+        image_channels = 3
+        if self.autoencoder is not None:
+            image_size = image_size // self.autoenc_scale_reduction
+            image_channels = self.autoenc_latent_channels
+        return jax.random.normal(rngs, (num_images, image_size, image_size, image_channels)) * variance
 
     def generate_images(self,
                         num_images=16, 
@@ -75,18 +120,23 @@ def generate_images(self,
                         end_step:int = 0,
                         steps_override=None,
                         priors=None, 
-                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42))) -> jnp.ndarray:
+                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42)),
+                        model_conditioning_inputs:tuple=()
+                        ) -> 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")
+            if self.autoencoder is not None:
+                priors = self.autoencoder.encode(priors)
             samples = priors
 
-        @jax.jit
+        # @jax.jit
         def sample_step(state:RandomMarkovState, samples, current_step, next_step):
             samples, state = self.sample_step(current_samples=samples,
                                               current_step=current_step, 
+                                              model_conditioning_inputs=model_conditioning_inputs,
                                               state=state, next_step=next_step)
             return samples, state
 
@@ -108,6 +158,8 @@ def sample_step(state:RandomMarkovState, 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, _, _ = self.sample_model(samples, current_step * step_ones, *model_conditioning_inputs)
+        if self.autoencoder is not None:
+            samples = self.autoencoder.decode(samples)
         samples = clip_images(samples)
-        return samples
+        return samples

flaxdiff/samplers/ddim.py

@@ -1,10 +1,10 @@
 import jax.numpy as jnp
 from .common import DiffusionSampler
-from ..utils import MarkovState
+from ..utils import MarkovState, RandomMarkovState
 
 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]:
+    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
+        return reconstructed_samples * next_signal_rate + pred_noise * next_noise_rate, state
+    

flaxdiff/samplers/ddpm.py

@@ -3,9 +3,8 @@
 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]:
+    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)
 
@@ -19,9 +18,8 @@ def generate_images(self, num_images=16, diffusion_steps=1000, start_step: int =
         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]:
+    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)
 
@@ -33,11 +31,7 @@ def take_next_step(self,
 
         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)
+        gamma = jnp.sqrt(noise_ratio_squared * (1 - signal_ratio_squared))
 
         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

@@ -5,9 +5,8 @@
 
 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]:
+    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)
 
@@ -22,9 +21,8 @@ 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]:
+    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)
 
@@ -38,9 +36,8 @@ 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]:
+    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)
 
@@ -56,4 +53,4 @@ def take_next_step(self,
         dW = jax.random.normal(subkey, current_samples.shape) * sigma_up
 
         next_samples = current_samples + dx * dt + dW
-        return next_samples, state
+        return next_samples, state

flaxdiff/samplers/heun_sampler.py

@@ -4,9 +4,8 @@
 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]:
+    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)
@@ -18,7 +17,7 @@ def take_next_step(self,
         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)
+        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

flaxdiff/samplers/multistep_dpm.py

@@ -8,9 +8,8 @@ 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]:
+    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)