flaxdiff 0.1.35.5__py3-none-any.whl → 0.1.36__py3-none-any.whl

flaxdiff/samplers/ddpm.py

@@ -3,9 +3,8 @@ 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]:
+    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 @@ class DDPMSampler(DiffusionSampler):
         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 @@ class SimpleDDPMSampler(DiffusionSampler):
         
         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 @@ 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]:
+    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 @@ class EulerAncestralSampler(DiffusionSampler):
         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 .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]:
+    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 @@ 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)
+        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 @@ class MultiStepDPM(DiffusionSampler):
         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)

flaxdiff/samplers/rk4_sampler.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 from .common import DiffusionSampler
-from ..utils import RandomMarkovState
+from ..utils import RandomMarkovState, MarkovState
 from ..schedulers import GeneralizedNoiseScheduler
 
 class RK4Sampler(DiffusionSampler):
@@ -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) -> tuple[jnp.ndarray, RandomMarkovState]:
+        def get_derivative(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)
+            x_0, eps, _ = self.sample_model(x_t, t, *model_conditioning_inputs)
             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]:
+    def sample_step(self, 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)
-        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)
+        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)
 
-        next_samples = current_samples + ((k1 + 2 * k2 + 2 * k3 + k4) / 6) * dt
+        next_samples = current_samples + (((k1 + 2 * k2 + 2 * k3 + k4) * dt) / 6)
         return next_samples, state

flaxdiff/trainer/autoencoder_trainer.py

@@ -14,7 +14,7 @@ from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransfor
 from flaxdiff.utils import RandomMarkovState
 
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
-
+from .diffusion_trainer import TrainState
 from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
 
 class AutoEncoderTrainer(SimpleTrainer):

flaxdiff/trainer/diffusion_trainer.py

@@ -1,22 +1,27 @@
+import flax
 from flax import linen as nn
 import jax
 from typing import Callable
 from dataclasses import field
 import jax.numpy as jnp
+import traceback
 import optax
+import functools
 from jax.sharding import Mesh, PartitionSpec as P
 from jax.experimental.shard_map import shard_map
-from typing import Dict, Callable, Sequence, Any, Union, Tuple
+from typing import Dict, Callable, Sequence, Any, Union, Tuple, Type
 
 from ..schedulers import NoiseScheduler
 from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+from ..samplers.common import DiffusionSampler
 
 from flaxdiff.utils import RandomMarkovState
 
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
 
 from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
-from flax.training.dynamic_scale import DynamicScale
+from flax.training import dynamic_scale as dynamic_scale_lib
+from flaxdiff.utils import TextEncoder, ConditioningEncoder
 
 class TrainState(SimpleTrainState):
     rngs: jax.random.PRNGKey
@@ -47,6 +52,7 @@ class DiffusionTrainer(SimpleTrainer):
                  name: str = "Diffusion",
                  model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
                  autoencoder: AutoEncoder = None,
+                 encoder: ConditioningEncoder = None,
                  **kwargs
                  ):
         super().__init__(
@@ -62,6 +68,7 @@ class DiffusionTrainer(SimpleTrainer):
         self.unconditional_prob = unconditional_prob
         
         self.autoencoder = autoencoder
+        self.encoder = encoder
 
     def generate_states(
         self,
@@ -84,8 +91,7 @@ class DiffusionTrainer(SimpleTrainer):
             new_state = existing_state
 
         if param_transforms is not None:
-            new_state['params'] = param_transforms(new_state['params'])
-            new_state['ema_params'] = param_transforms(new_state['ema_params'])
+            params = param_transforms(params)
 
         state = TrainState.create(
             apply_fn=model.apply,
@@ -94,7 +100,7 @@ class DiffusionTrainer(SimpleTrainer):
             tx=optimizer,
             rngs=rngs,
             metrics=Metrics.empty(),
-            dynamic_scale = DynamicScale() if use_dynamic_scale else None
+            dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None
         )
             
         if existing_best_state is not None:
@@ -105,7 +111,7 @@ class DiffusionTrainer(SimpleTrainer):
 
         return state, best_state
 
-    def _define_train_step(self, batch_size, null_labels_seq, text_embedder):
+    def _define_train_step(self, batch_size):
         noise_schedule: NoiseScheduler = self.noise_schedule
         model = self.model
         model_output_transform = self.model_output_transform
@@ -114,6 +120,11 @@ class DiffusionTrainer(SimpleTrainer):
 
         # Determine the number of unconditional samples
         num_unconditional = int(batch_size * unconditional_prob)
+        
+        null_labels_full = self.encoder([""])
+        null_labels_seq = jnp.array(null_labels_full[0], dtype=jnp.float16)
+        
+        conditioning_encoder = self.encoder
 
         nS, nC = null_labels_seq.shape
         null_labels_seq = jnp.broadcast_to(
@@ -131,6 +142,11 @@ class DiffusionTrainer(SimpleTrainer):
             local_rng_state = RandomMarkovState(subkey)
             
             images = batch['image']
+            
+            # First get the standard deviation of the images
+            # std = jnp.std(images, axis=(1, 2, 3))
+            # is_non_zero = (std > 0)
+            
             images = jnp.array(images, dtype=jnp.float32)
             # normalize image
             images = (images - 127.5) / 127.5
@@ -140,9 +156,7 @@ class DiffusionTrainer(SimpleTrainer):
                 local_rng_state, rngs = local_rng_state.get_random_key()
                 images = autoencoder.encode(images, rngs)
 
-            output = text_embedder(
-                input_ids=batch['input_ids'], attention_mask=batch['attention_mask'])
-            label_seq = output.last_hidden_state
+            label_seq = conditioning_encoder.encode_from_tokens(batch)
 
             # Generate random probabilities to decide how much of this batch will be unconditional
 
@@ -163,8 +177,11 @@ class DiffusionTrainer(SimpleTrainer):
                 preds = model_output_transform.pred_transform(
                     noisy_images, preds, rates)
                 nloss = loss_fn(preds, expected_output)
-                # nloss = jnp.mean(nloss, axis=1)
+                # Ignore the loss contribution of images with zero standard deviation
                 nloss *= noise_schedule.get_weights(noise_level)
+                # nloss = jnp.mean(nloss, axis=(1,2,3))
+                # nloss = jnp.where(is_non_zero, nloss, 0)
+                # nloss = jnp.mean(nloss, where=nloss != 0)
                 nloss = jnp.mean(nloss)
                 loss = nloss
                 return loss
@@ -185,11 +202,11 @@ class DiffusionTrainer(SimpleTrainer):
             
             new_state = train_state.apply_gradients(grads=grads)
             
-            if train_state.dynamic_scale:
+            if train_state.dynamic_scale is not None:
                 # if is_fin == False the gradients contain Inf/NaNs and optimizer state and
                 # params should be restored (= skip this step).
                 select_fn = functools.partial(jnp.where, is_fin)
-                new_state = train_state.replace(
+                new_state = new_state.replace(
                     opt_state=jax.tree_util.tree_map(
                         select_fn, new_state.opt_state, train_state.opt_state
                     ),
@@ -211,24 +228,99 @@ class DiffusionTrainer(SimpleTrainer):
             
         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):
-        null_labels_full = data['null_labels_full']
+    def _define_vaidation_step(self, sampler_class: Type[DiffusionSampler]):
+        model = self.model
+        encoder = self.encoder
+        autoencoder = self.autoencoder
+        
+        null_labels_full = encoder([""])
+        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,
+                model_output_transform=self.model_output_transform,
+                image_size=self.input_shapes['x'][0],
+                null_labels_seq=null_labels_full,
+                autoencoder=autoencoder,
+            )
+            return sampler
+        
+        def generate_samples(
+            batch,
+            sampler: DiffusionSampler, 
+            diffusion_steps: int,
+        ):
+            labels_seq = encoder.encode_from_tokens(batch)
+            labels_seq = jnp.array(labels_seq, dtype=jnp.float16)
+            samples = sampler.generate_images(
+                num_images=len(labels_seq),
+                diffusion_steps=diffusion_steps,
+                start_step=1000,
+                end_step=0,
+                priors=None,
+                model_conditioning_inputs=(labels_seq,),
+            )
+            return samples
+        
+        return generate_sampler, generate_samples
+
+    def validation_loop(
+        self,
+        val_state: SimpleTrainState,
+        val_step_fn: Callable,
+        val_ds,
+        val_steps_per_epoch,
+        current_step,
+        diffusion_steps=200,
+    ):
+        generate_sampler, generate_samples = val_step_fn
+        
+        sampler = generate_sampler(val_state)
+        
+        val_ds = iter(val_ds()) if val_ds else None
+        # Evaluation step
+        try:
+            samples = generate_samples(
+                next(val_ds),
+                sampler,
+                diffusion_steps,
+            )
+            
+            # Put each sample on wandb
+            if self.wandb:
+                import numpy as np
+                from wandb import Image as wandbImage
+                wandb_images = []
+                for i in range(samples.shape[0]):
+                    # convert the sample to numpy
+                    sample = np.array(samples[i])
+                    # denormalize the image
+                    sample = (sample + 1) * 127.5
+                    sample = np.clip(sample, 0, 255).astype(np.uint8)
+                    # add the image to the list
+                    wandb_images.append(sample)
+                    # log the images to wandb
+                    self.wandb.log({
+                        f"sample_{i}": wandbImage(sample, caption=f"Sample {i} at step {current_step}")
+                    }, step=current_step)
+        except Exception as e:
+            print("Error logging images to wandb", e)
+            traceback.print_exc()
+    
+    def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class=None):
         local_batch_size = data['local_batch_size']
-        text_embedder = data['model']
-        super().fit(data, steps_per_epoch, epochs, {
-            "batch_size": local_batch_size, "null_labels_seq": null_labels_full, "text_embedder": text_embedder})
-
-def boolean_string(s):
-    if type(s) == bool:
-        return s
-    return s == 'True'
+        validation_step_args = {
+            "sampler_class": sampler_class,
+        }
+        super().fit(
+            data, 
+            train_steps_per_epoch=training_steps_per_epoch, 
+            epochs=epochs, 
+            train_step_args={"batch_size": local_batch_size}, 
+            val_steps_per_epoch=val_steps_per_epoch,
+            validation_step_args=validation_step_args,
+        )