feat: stability improvements

Author: AshishKumar4

flaxdiff/predictors/__init__.py

@@ -81,16 +81,16 @@ def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray])
         epsilon = (x_t - x_0 * signal_rate) / noise_rate
         return x_0, epsilon
 
-    def pred_transform(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray:
+    def pred_transform(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) -> jnp.ndarray:
         _, sigma = rates
-        c_out = sigma * self.sigma_data / jnp.sqrt(self.sigma_data ** 2 + sigma ** 2)
-        c_skip = self.sigma_data ** 2 / (self.sigma_data ** 2 + sigma ** 2)
+        c_out = sigma * self.sigma_data / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon)
+        c_skip = self.sigma_data ** 2 / (self.sigma_data ** 2 + sigma ** 2 + epsilon)
         c_out = c_out.reshape((-1, 1, 1, 1))
         c_skip = c_skip.reshape((-1, 1, 1, 1))
         x_0 = c_out * preds + c_skip * x_t
         return x_0
 
-    def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray:
+    def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) -> jnp.ndarray:
         _, sigma = rates
-        c_in = 1 / jnp.sqrt(self.sigma_data ** 2 + sigma ** 2)
+        c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon)
         return c_in

flaxdiff/trainer/diffusion_trainer.py

@@ -167,7 +167,10 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc
             noise_level, local_rng_state = noise_schedule.generate_timesteps(images.shape[0], local_rng_state)
 
             local_rng_state, rngs = local_rng_state.get_random_key()
-            noise: jax.Array = jax.random.normal(rngs, shape=images.shape)
+            noise: jax.Array = jax.random.normal(rngs, shape=images.shape, dtype=jnp.float32)
+            
+            # Make sure image is also float32
+            images = images.astype(jnp.float32)
 
             rates = noise_schedule.get_rates(noise_level)
             noisy_images, c_in, expected_output = model_output_transform.forward_diffusion(
@@ -197,8 +200,23 @@ def model_loss(params):
                 loss, grads = grad_fn(train_state.params)
                 if distributed_training:
                     grads = jax.lax.pmean(grads, "data")
+                    
+            # # check gradients for NaN/Inf
+            # has_nan_or_inf = jax.tree_util.tree_reduce(
+            #     lambda acc, x: jnp.logical_or(acc, jnp.logical_or(jnp.isnan(x).any(), jnp.isinf(x).any())),
+            #     grads,
+            #     initializer=False
+            # )
 
-            new_state = train_state.apply_gradients(grads=grads)
+            # # Only apply gradients if they're valid
+            # new_state = jax.lax.cond(
+            #     has_nan_or_inf,
+            #     lambda _: train_state,  # Skip gradient update
+            #     lambda _: train_state.apply_gradients(grads=grads),
+            #     operand=None
+            # )
+    
+            # new_state = train_state.apply_gradients(grads=grads)
 
             if train_state.dynamic_scale is not None:
                 # if is_fin == False the gradients contain Inf/NaNs and optimizer state and

flaxdiff/trainer/simple_trainer.py

@@ -403,7 +403,6 @@ def train_loop(
         rng_state
     ):
         global_device_count = jax.device_count()
-        local_device_count = jax.local_device_count()
         process_index = jax.process_index()
         if self.distributed_training:
             global_device_indexes = jnp.arange(global_device_count)
@@ -434,11 +433,16 @@ def train_loop(
                 # loss = jax.experimental.multihost_utils.process_allgather(loss)
                 loss = jnp.mean(loss) # Just to make sure its a scaler value
 
-            if loss <= 1e-6:
+            if loss <= 1e-8:
                 # If the loss is too low, we can assume the model has diverged
                 print(colored(f"Loss too low at step {current_step} => {loss}", 'red'))
                 # Reset the model to the old state
-                exit(1)
+                if self.best_state is not None:
+                    print(colored(f"Resetting model to best state", 'red'))
+                    train_state = self.best_state
+                    loss = self.best_loss
+                else:
+                    exit(1)
 
             epoch_loss += loss
             current_step += 1

pyproject.toml

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