flaxdiff 0.1.36.4__tar.gz → 0.1.37__tar.gz

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36.4
+Version: 0.1.37
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT
@@ -96,7 +96,7 @@ Also, few of the text may be generated with help of github copilot, so please ex
 ### Schedulers
 Implemented in `flaxdiff.schedulers`:
 - **LinearNoiseSchedule** (`flaxdiff.schedulers.LinearNoiseSchedule`): A beta-parameterized discrete scheduler.
-- **CosineNoiseSchedule** (`flaxdiff.schedulers.CosineNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **CosineNoiseScheduler** (`flaxdiff.schedulers.CosineNoiseScheduler`): A beta-parameterized discrete scheduler.
 - **ExpNoiseSchedule** (`flaxdiff.schedulers.ExpNoiseSchedule`): A beta-parameterized discrete scheduler.
 - **CosineContinuousNoiseScheduler** (`flaxdiff.schedulers.CosineContinuousNoiseScheduler`): A continuous scheduler.
 - **CosineGeneralNoiseScheduler** (`flaxdiff.schedulers.CosineGeneralNoiseScheduler`): A continuous sigma parameterized cosine scheduler.
@@ -147,43 +147,81 @@ sticking to the versions mentioned in the requirements.txt
 Here is a simplified example to get you started with training a diffusion model using FlaxDiff:
 
 ```python
-from flaxdiff.schedulers import EDMNoiseScheduler
+from flaxdiff.schedulers import EDMNoiseScheduler, KarrasVENoiseScheduler
 from flaxdiff.predictors import KarrasPredictionTransform
-from flaxdiff.models.simple_unet import SimpleUNet as UNet
+from flaxdiff.models.simple_unet import Unet
 from flaxdiff.trainer import DiffusionTrainer
+from flaxdiff.data.datasets import get_dataset_grain
+from flaxdiff.utils import defaultTextEncodeModel
+from flaxdiff.samplers.euler import EulerAncestralSampler
 import jax
+import jax.numpy as jnp
 import optax
 from datetime import datetime
 
 BATCH_SIZE = 16
-IMAGE_SIZE = 64
+IMAGE_SIZE = 128
 
 # Define noise scheduler
 edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
-
+karas_ve_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
 # Define model
-unet = UNet(emb_features=256, 
-            feature_depths=[64, 128, 256, 512],
-            attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}],
+unet = Unet(emb_features=256, 
+            feature_depths=[64, 64, 128, 256, 512],
+            attention_configs=[
+                None,
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":False, "use_self_and_cross":False}
+                ],
             num_res_blocks=2,
-            num_middle_res_blocks=1)
-
+            num_middle_res_blocks=1
+)
 # Load dataset
-data, datalen = get_dataset("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+data = get_dataset_grain("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+datalen = data['train_len']
 batches = datalen // BATCH_SIZE
 
+input_shapes = {
+    "x": (IMAGE_SIZE, IMAGE_SIZE, 3),
+    "temb": (),
+    "textcontext": (77, 768)
+}
+text_encoder = defaultTextEncodeModel()
+
+# Construct a validation set by the prompts
+val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold']
+
+def get_val_dataset(batch_size=8):
+    for i in range(0, len(val_prompts), batch_size):
+        prompts = val_prompts[i:i + batch_size]
+        tokens = text_encoder.tokenize(prompts)
+        yield tokens
+
+data['test'] = get_val_dataset
+data['test_len'] = len(val_prompts)
+
 # Define optimizer
 solver = optax.adam(2e-4)
 
 # Create trainer
-trainer = DiffusionTrainer(unet, optimizer=solver, 
-                           noise_schedule=edm_schedule,
-                           rngs=jax.random.PRNGKey(4), 
-                           name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
-                           model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data))
+trainer = DiffusionTrainer(
+    unet, optimizer=solver, 
+    input_shapes=input_shapes,
+    noise_schedule=edm_schedule,
+    rngs=jax.random.PRNGKey(4), 
+    name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
+    model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data),
+    encoder=text_encoder,
+    distributed_training=True,
+    wandb_config = {
+        "project": 'mlops-msml605-project',
+        "name": f"prototype-{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}",
+})
 
 # Train the model
-final_state = trainer.fit(data, batches, epochs=2000)
+final_state = trainer.fit(data, batches, epochs=2000, sampler_class=EulerAncestralSampler, sampling_noise_schedule=karas_ve_schedule)
 ```
 
 ### Inference Example
@@ -301,8 +339,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],
@@ -321,8 +359,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/README.md

@@ -74,7 +74,7 @@ Also, few of the text may be generated with help of github copilot, so please ex
 ### Schedulers
 Implemented in `flaxdiff.schedulers`:
 - **LinearNoiseSchedule** (`flaxdiff.schedulers.LinearNoiseSchedule`): A beta-parameterized discrete scheduler.
-- **CosineNoiseSchedule** (`flaxdiff.schedulers.CosineNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **CosineNoiseScheduler** (`flaxdiff.schedulers.CosineNoiseScheduler`): A beta-parameterized discrete scheduler.
 - **ExpNoiseSchedule** (`flaxdiff.schedulers.ExpNoiseSchedule`): A beta-parameterized discrete scheduler.
 - **CosineContinuousNoiseScheduler** (`flaxdiff.schedulers.CosineContinuousNoiseScheduler`): A continuous scheduler.
 - **CosineGeneralNoiseScheduler** (`flaxdiff.schedulers.CosineGeneralNoiseScheduler`): A continuous sigma parameterized cosine scheduler.
@@ -125,43 +125,81 @@ sticking to the versions mentioned in the requirements.txt
 Here is a simplified example to get you started with training a diffusion model using FlaxDiff:
 
 ```python
-from flaxdiff.schedulers import EDMNoiseScheduler
+from flaxdiff.schedulers import EDMNoiseScheduler, KarrasVENoiseScheduler
 from flaxdiff.predictors import KarrasPredictionTransform
-from flaxdiff.models.simple_unet import SimpleUNet as UNet
+from flaxdiff.models.simple_unet import Unet
 from flaxdiff.trainer import DiffusionTrainer
+from flaxdiff.data.datasets import get_dataset_grain
+from flaxdiff.utils import defaultTextEncodeModel
+from flaxdiff.samplers.euler import EulerAncestralSampler
 import jax
+import jax.numpy as jnp
 import optax
 from datetime import datetime
 
 BATCH_SIZE = 16
-IMAGE_SIZE = 64
+IMAGE_SIZE = 128
 
 # Define noise scheduler
 edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
-
+karas_ve_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
 # Define model
-unet = UNet(emb_features=256, 
-            feature_depths=[64, 128, 256, 512],
-            attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}],
+unet = Unet(emb_features=256, 
+            feature_depths=[64, 64, 128, 256, 512],
+            attention_configs=[
+                None,
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":False, "use_self_and_cross":False}
+                ],
             num_res_blocks=2,
-            num_middle_res_blocks=1)
-
+            num_middle_res_blocks=1
+)
 # Load dataset
-data, datalen = get_dataset("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+data = get_dataset_grain("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+datalen = data['train_len']
 batches = datalen // BATCH_SIZE
 
+input_shapes = {
+    "x": (IMAGE_SIZE, IMAGE_SIZE, 3),
+    "temb": (),
+    "textcontext": (77, 768)
+}
+text_encoder = defaultTextEncodeModel()
+
+# Construct a validation set by the prompts
+val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold']
+
+def get_val_dataset(batch_size=8):
+    for i in range(0, len(val_prompts), batch_size):
+        prompts = val_prompts[i:i + batch_size]
+        tokens = text_encoder.tokenize(prompts)
+        yield tokens
+
+data['test'] = get_val_dataset
+data['test_len'] = len(val_prompts)
+
 # Define optimizer
 solver = optax.adam(2e-4)
 
 # Create trainer
-trainer = DiffusionTrainer(unet, optimizer=solver, 
-                           noise_schedule=edm_schedule,
-                           rngs=jax.random.PRNGKey(4), 
-                           name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
-                           model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data))
+trainer = DiffusionTrainer(
+    unet, optimizer=solver, 
+    input_shapes=input_shapes,
+    noise_schedule=edm_schedule,
+    rngs=jax.random.PRNGKey(4), 
+    name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
+    model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data),
+    encoder=text_encoder,
+    distributed_training=True,
+    wandb_config = {
+        "project": 'mlops-msml605-project',
+        "name": f"prototype-{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}",
+})
 
 # Train the model
-final_state = trainer.fit(data, batches, epochs=2000)
+final_state = trainer.fit(data, batches, epochs=2000, sampler_class=EulerAncestralSampler, sampling_noise_schedule=karas_ve_schedule)
 ```
 
 ### Inference Example
@@ -279,8 +317,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],
@@ -299,8 +337,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/common.py

@@ -37,7 +37,7 @@ class DiffusionSampler():
             # 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):
+            def sample_model(params, 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)
@@ -46,7 +46,7 @@ class DiffusionSampler():
                 
                 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)
+                model_output = self.model.apply(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)
@@ -55,10 +55,10 @@ class DiffusionSampler():
                 return x_0, eps, model_output
         else:
             # Unconditional sampling
-            def sample_model(x_t, t, *additional_inputs):
+            def sample_model(params, 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)
+                model_output = self.model.apply(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
             
@@ -70,22 +70,23 @@ class DiffusionSampler():
         self.sample_model = sample_model
 
     # Used to sample from the diffusion model
-    def sample_step(self, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
+    def sample_step(self, sample_model_fn, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
         # First clip the noisy 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, *model_conditioning_inputs)
+        pred_images, pred_noise, _ = sample_model_fn(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,
-                             model_conditioning_inputs=model_conditioning_inputs
+                                pred_noise=pred_noise, current_step=current_step, next_step=next_step, state=state,
+                                model_conditioning_inputs=model_conditioning_inputs,
+                                sample_model_fn=sample_model_fn,
                              )
         return new_samples, state
 
     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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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
@@ -114,6 +115,7 @@ class DiffusionSampler():
         return jax.random.normal(rngs, (num_images, image_size, image_size, image_channels)) * variance
 
     def generate_images(self,
+                        params:dict=None,
                         num_images=16, 
                         diffusion_steps=1000, 
                         start_step:int = None,
@@ -131,10 +133,15 @@ class DiffusionSampler():
             if self.autoencoder is not None:
                 priors = self.autoencoder.encode(priors)
             samples = priors
+            
+        params = params if params is not None else self.params
+        
+        def sample_model_fn(x_t, t, *additional_inputs):
+            return self.sample_model(params, x_t, t, *additional_inputs)
 
         # @jax.jit
-        def sample_step(state:RandomMarkovState, samples, current_step, next_step):
-            samples, state = self.sample_step(current_samples=samples,
+        def sample_step(sample_model_fn, state:RandomMarkovState, samples, current_step, next_step):
+            samples, state = self.sample_step(sample_model_fn=sample_model_fn, current_samples=samples,
                                               current_step=current_step, 
                                               model_conditioning_inputs=model_conditioning_inputs,
                                               state=state, next_step=next_step)
@@ -154,11 +161,11 @@ class DiffusionSampler():
             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)
+                samples, rngstate = sample_step(sample_model_fn, 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, *model_conditioning_inputs)
+                samples, _, _ = sample_model_fn(samples, current_step * step_ones, *model_conditioning_inputs)
         if self.autoencoder is not None:
             samples = self.autoencoder.decode(samples)
         samples = clip_images(samples)

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/ddim.py

@@ -4,7 +4,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/ddpm.py

@@ -4,7 +4,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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,7 +19,7 @@ class DDPMSampler(DiffusionSampler):
     
 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         state, rng = state.get_random_key()
         noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32)
 

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/euler.py

@@ -6,7 +6,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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,7 +22,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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)
 
@@ -37,7 +37,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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)
 

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/heun_sampler.py

@@ -5,7 +5,7 @@ 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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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)
@@ -17,7 +17,7 @@ class HeunSampler(DiffusionSampler):
         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)
+        estimated_x_0, _, _ = sample_model_fn(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-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/multistep_dpm.py

@@ -9,7 +9,7 @@ class MultiStepDPM(DiffusionSampler):
         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]:
+                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, 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)

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/samplers/rk4_sampler.py

@@ -9,14 +9,14 @@ class RK4Sampler(DiffusionSampler):
         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]:
+        def get_derivative(sample_model_fn, 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)
+            x_0, eps, _ = sample_model_fn(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]:
+    def sample_step(self, sample_model_fn, 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
@@ -25,10 +25,10 @@ class RK4Sampler(DiffusionSampler):
 
         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)
+        k1, state = self.get_derivative(sample_model_fn, current_samples, current_sigma, state, model_conditioning_inputs)
+        k2, state = self.get_derivative(sample_model_fn, current_samples + 0.5 * k1 * dt, current_sigma + 0.5 * dt, state, model_conditioning_inputs)
+        k3, state = self.get_derivative(sample_model_fn, current_samples + 0.5 * k2 * dt, current_sigma + 0.5 * dt, state, model_conditioning_inputs)
+        k4, state = self.get_derivative(sample_model_fn, 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-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/schedulers/__init__.py

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

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/schedulers/cosine.py

@@ -12,7 +12,7 @@ def cosine_beta_schedule(timesteps, start_angle=0.008, end_angle=0.999):
     betas = 1 - (alphas_bar[1:] / alphas_bar[:-1])
     return np.clip(betas, 0, end_angle)
 
-class CosineNoiseSchedule(DiscreteNoiseScheduler):
+class CosineNoiseScheduler(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)
 

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff/trainer/diffusion_trainer.py

@@ -235,19 +235,19 @@ class DiffusionTrainer(SimpleTrainer):
         null_labels_full = null_labels_full.astype(jnp.float16)
         # null_labels_seq = jnp.array(null_labels_full[0], dtype=jnp.float16)
         
-        def generate_sampler(state: TrainState):
-            sampler = sampler_class(
-                model=model,
-                params=state.ema_params,
-                noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule,
-                model_output_transform=self.model_output_transform,
-                image_size=self.input_shapes['x'][0],
-                null_labels_seq=null_labels_full,
-                autoencoder=autoencoder,
-            )
-            return sampler
+        sampler = sampler_class(
+            model=model,
+            params=None,
+            noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule,
+            model_output_transform=self.model_output_transform,
+            image_size=self.input_shapes['x'][0],
+            null_labels_seq=null_labels_full,
+            autoencoder=autoencoder,
+            guidance_scale=3.0,
+        )
         
         def generate_samples(
+            val_state: TrainState,
             batch,
             sampler: DiffusionSampler, 
             diffusion_steps: int,
@@ -255,6 +255,7 @@ class DiffusionTrainer(SimpleTrainer):
             labels_seq = encoder.encode_from_tokens(batch)
             labels_seq = jnp.array(labels_seq, dtype=jnp.float16)
             samples = sampler.generate_images(
+                params=val_state.ema_params,
                 num_images=len(labels_seq),
                 diffusion_steps=diffusion_steps,
                 start_step=1000,
@@ -264,7 +265,7 @@ class DiffusionTrainer(SimpleTrainer):
             )
             return samples
         
-        return generate_sampler, generate_samples
+        return sampler, generate_samples
 
     def validation_loop(
         self,
@@ -275,14 +276,15 @@ class DiffusionTrainer(SimpleTrainer):
         current_step,
         diffusion_steps=200,
     ):
-        generate_sampler, generate_samples = val_step_fn
+        sampler, generate_samples = val_step_fn
         
-        sampler = generate_sampler(val_state)
+        # sampler = generate_sampler(val_state)
         
         val_ds = iter(val_ds()) if val_ds else None
         # Evaluation step
         try:
             samples = generate_samples(
+                val_state,
                 next(val_ds),
                 sampler,
                 diffusion_steps,

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36.4
+Version: 0.1.37
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT
@@ -96,7 +96,7 @@ Also, few of the text may be generated with help of github copilot, so please ex
 ### Schedulers
 Implemented in `flaxdiff.schedulers`:
 - **LinearNoiseSchedule** (`flaxdiff.schedulers.LinearNoiseSchedule`): A beta-parameterized discrete scheduler.
-- **CosineNoiseSchedule** (`flaxdiff.schedulers.CosineNoiseSchedule`): A beta-parameterized discrete scheduler.
+- **CosineNoiseScheduler** (`flaxdiff.schedulers.CosineNoiseScheduler`): A beta-parameterized discrete scheduler.
 - **ExpNoiseSchedule** (`flaxdiff.schedulers.ExpNoiseSchedule`): A beta-parameterized discrete scheduler.
 - **CosineContinuousNoiseScheduler** (`flaxdiff.schedulers.CosineContinuousNoiseScheduler`): A continuous scheduler.
 - **CosineGeneralNoiseScheduler** (`flaxdiff.schedulers.CosineGeneralNoiseScheduler`): A continuous sigma parameterized cosine scheduler.
@@ -147,43 +147,81 @@ sticking to the versions mentioned in the requirements.txt
 Here is a simplified example to get you started with training a diffusion model using FlaxDiff:
 
 ```python
-from flaxdiff.schedulers import EDMNoiseScheduler
+from flaxdiff.schedulers import EDMNoiseScheduler, KarrasVENoiseScheduler
 from flaxdiff.predictors import KarrasPredictionTransform
-from flaxdiff.models.simple_unet import SimpleUNet as UNet
+from flaxdiff.models.simple_unet import Unet
 from flaxdiff.trainer import DiffusionTrainer
+from flaxdiff.data.datasets import get_dataset_grain
+from flaxdiff.utils import defaultTextEncodeModel
+from flaxdiff.samplers.euler import EulerAncestralSampler
 import jax
+import jax.numpy as jnp
 import optax
 from datetime import datetime
 
 BATCH_SIZE = 16
-IMAGE_SIZE = 64
+IMAGE_SIZE = 128
 
 # Define noise scheduler
 edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
-
+karas_ve_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
 # Define model
-unet = UNet(emb_features=256, 
-            feature_depths=[64, 128, 256, 512],
-            attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}],
+unet = Unet(emb_features=256, 
+            feature_depths=[64, 64, 128, 256, 512],
+            attention_configs=[
+                None,
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":True, "use_self_and_cross":True}, 
+                {"heads":8, "dtype":jnp.float16, "flash_attention":False, "use_projection":False, "use_self_and_cross":False}
+                ],
             num_res_blocks=2,
-            num_middle_res_blocks=1)
-
+            num_middle_res_blocks=1
+)
 # Load dataset
-data, datalen = get_dataset("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+data = get_dataset_grain("oxford_flowers102", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)
+datalen = data['train_len']
 batches = datalen // BATCH_SIZE
 
+input_shapes = {
+    "x": (IMAGE_SIZE, IMAGE_SIZE, 3),
+    "temb": (),
+    "textcontext": (77, 768)
+}
+text_encoder = defaultTextEncodeModel()
+
+# Construct a validation set by the prompts
+val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold']
+
+def get_val_dataset(batch_size=8):
+    for i in range(0, len(val_prompts), batch_size):
+        prompts = val_prompts[i:i + batch_size]
+        tokens = text_encoder.tokenize(prompts)
+        yield tokens
+
+data['test'] = get_val_dataset
+data['test_len'] = len(val_prompts)
+
 # Define optimizer
 solver = optax.adam(2e-4)
 
 # Create trainer
-trainer = DiffusionTrainer(unet, optimizer=solver, 
-                           noise_schedule=edm_schedule,
-                           rngs=jax.random.PRNGKey(4), 
-                           name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
-                           model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data))
+trainer = DiffusionTrainer(
+    unet, optimizer=solver, 
+    input_shapes=input_shapes,
+    noise_schedule=edm_schedule,
+    rngs=jax.random.PRNGKey(4), 
+    name="Diffusion_SDE_VE_" + datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
+    model_output_transform=KarrasPredictionTransform(sigma_data=edm_schedule.sigma_data),
+    encoder=text_encoder,
+    distributed_training=True,
+    wandb_config = {
+        "project": 'mlops-msml605-project',
+        "name": f"prototype-{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}",
+})
 
 # Train the model
-final_state = trainer.fit(data, batches, epochs=2000)
+final_state = trainer.fit(data, batches, epochs=2000, sampler_class=EulerAncestralSampler, sampling_noise_schedule=karas_ve_schedule)
 ```
 
 ### Inference Example
@@ -301,8 +339,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],
@@ -321,8 +359,8 @@ Images generated by the following prompts using classifier free guidance with gu
 `Training Epochs: 1000`
 `Steps per epoch: 511`
 
-`Training Noise Schedule: CosineNoiseSchedule`
-`Inference Noise Schedule: CosineNoiseSchedule`
+`Training Noise Schedule: CosineNoiseScheduler`
+`Inference Noise Schedule: CosineNoiseScheduler`
 
 `Model: UNet(emb_features=256, 
             feature_depths=[64, 128, 256, 512],

{flaxdiff-0.1.36.4 → flaxdiff-0.1.37}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.1.36.4"
+version = "0.1.37"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [