flaxdiff 0.1.36.5__tar.gz → 0.1.37.1__tar.gz

{flaxdiff-0.1.36.5 → flaxdiff-0.1.37.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36.5
+Version: 0.1.37.1
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

flaxdiff-0.1.37.1/flaxdiff/samplers/__init__.py

@@ -0,0 +1,7 @@
+from .common import *
+from .ddim import *
+from .ddpm import *
+from .euler import *
+from .heun_sampler import *
+from .rk4_sampler import *
+from .multistep_dpm import *

{flaxdiff-0.1.36.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/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.37.1/flaxdiff/schedulers/__init__.py

@@ -0,0 +1,6 @@
+from .discrete import *
+from .common import *
+from .cosine import *
+from .linear import *
+from .sqrt import *
+from .karras import *

{flaxdiff-0.1.36.5 → flaxdiff-0.1.37.1}/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.5 → flaxdiff-0.1.37.1}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36.5
+Version: 0.1.37.1
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.1.36.5 → flaxdiff-0.1.37.1}/pyproject.toml

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

flaxdiff-0.1.36.5/flaxdiff/samplers/__init__.py

@@ -1,7 +0,0 @@
-from .common import DiffusionSampler
-from .ddim import DDIMSampler
-from .ddpm import DDPMSampler, SimpleDDPMSampler
-from .euler import EulerSampler, SimplifiedEulerSampler
-from .heun_sampler import HeunSampler
-from .rk4_sampler import RK4Sampler
-from .multistep_dpm import MultiStepDPM

flaxdiff-0.1.36.5/flaxdiff/schedulers/__init__.py

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