flaxdiff 0.1.5__tar.gz → 0.1.6__tar.gz

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.5
+Version: 0.1.6
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com
@@ -13,6 +13,8 @@ Requires-Dist: clu
 
 # ![](images/logo.jpeg "FlaxDiff")
 
+**This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
+
 ## A Versatile and simple Diffusion Library
 
 In recent years, diffusion and score-based multi-step models have revolutionized the generative AI domain. However, the latest research in this field has become highly math-intensive, making it challenging to understand how state-of-the-art diffusion models work and generate such impressive images. Replicating this research in code can be daunting.

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/README.md

@@ -1,5 +1,7 @@
 # ![](images/logo.jpeg "FlaxDiff")
 
+**This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
+
 ## A Versatile and simple Diffusion Library
 
 In recent years, diffusion and score-based multi-step models have revolutionized the generative AI domain. However, the latest research in this field has become highly math-intensive, making it challenging to understand how state-of-the-art diffusion models work and generate such impressive images. Replicating this research in code can be daunting.

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff/models/attention.py

@@ -5,7 +5,8 @@ Some Code ported from https://github.com/huggingface/diffusers/blob/main/src/dif
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from typing import Dict, Callable, Sequence, Any, Union
+from typing import Dict, Callable, Sequence, Any, Union, Tuple, Optional
+from flax.typing import Dtype, PrecisionLike
 import einops
 import functools
 import math
@@ -18,8 +19,8 @@ class EfficientAttention(nn.Module):
     query_dim: int
     heads: int = 4
     dim_head: int = 64
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGHEST
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
     use_bias: bool = True
     kernel_init: Callable = lambda : kernel_init(1.0)
 
@@ -109,8 +110,8 @@ class NormalAttention(nn.Module):
     query_dim: int
     heads: int = 4
     dim_head: int = 64
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGHEST
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
     use_bias: bool = True
     kernel_init: Callable = lambda : kernel_init(1.0)
 
@@ -166,8 +167,8 @@ class BasicTransformerBlock(nn.Module):
     query_dim: int
     heads: int = 4
     dim_head: int = 64
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGHEST
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
     use_bias: bool = True
     kernel_init: Callable = lambda : kernel_init(1.0)
     use_flash_attention:bool = False
@@ -286,8 +287,8 @@ class BasicTransformerBlock(nn.Module):
     query_dim: int
     heads: int = 4
     dim_head: int = 64
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGHEST
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
     use_bias: bool = True
     kernel_init: Callable = lambda : kernel_init(1.0)
     use_flash_attention:bool = False
@@ -346,8 +347,8 @@ class TransformerBlock(nn.Module):
     heads: int = 4
     dim_head: int = 32
     use_linear_attention: bool = True
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
     use_projection: bool = False
     use_flash_attention:bool = True
     use_self_and_cross:bool = False

flaxdiff-0.1.6/flaxdiff/models/autoencoder/__init__.py

@@ -0,0 +1,2 @@
+from .autoencoder import AutoEncoder
+from .diffusers import StableDiffusionVAE

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff/models/autoencoder/autoencoder.py

@@ -6,9 +6,14 @@ import einops
 from ..common import kernel_init, ConvLayer, Upsample, Downsample, PixelShuffle
 
 
-class AutoEncoder:
+class AutoEncoder():
     def encode(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
         raise NotImplementedError
     
     def decode(self, z: jnp.ndarray, **kwargs) -> jnp.ndarray:
-        raise NotImplementedError
+        raise NotImplementedError
+    
+    def __call__(self, x: jnp.ndarray):
+        latents = self.encode(x)
+        reconstructions = self.decode(latents)
+        return reconstructions

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff/models/autoencoder/diffusers.py

@@ -5,6 +5,9 @@ from .autoencoder import AutoEncoder
 
 """
 This module contains an Autoencoder implementation which uses the Stable Diffusion VAE model from the HuggingFace Diffusers library.
+The actual model was not trained by me, but was taken from the HuggingFace model hub.
+I have only implemented the wrapper around the diffusers pipeline to make it compatible with our library
+All credits for the model go to the developers of Stable Diffusion VAE and all credits for the pipeline go to the developers of the Diffusers library.
 """
 
 class StableDiffusionVAE(AutoEncoder):

flaxdiff-0.1.6/flaxdiff/models/autoencoder/simple_autoenc.py

@@ -0,0 +1,26 @@
+from typing import Any, List, Optional, Callable
+import jax
+import flax.linen as nn
+from jax import numpy as jnp
+from flax.typing import Dtype, PrecisionLike
+from .autoencoder import AutoEncoder
+
+class SimpleAutoEncoder(AutoEncoder):
+    latent_channels: int
+    feature_depths: List[int]=[64, 128, 256, 512]
+    attention_configs:list=[{"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}],
+    num_res_blocks: int=2
+    num_middle_res_blocks:int=1,
+    activation:Callable = jax.nn.swish
+    norm_groups:int=8
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    
+    # def encode(self, x: jnp.ndarray):
+        
+    
+    @nn.compact
+    def __call__(self, x: jnp.ndarray):
+        latents = self.encode(x)
+        reconstructions = self.decode(latents)
+        return reconstructions

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff/models/common.py

@@ -1,6 +1,8 @@
 import jax.numpy as jnp
 import jax
 from flax import linen as nn
+from typing import Optional, Any, Callable, Sequence, Union
+from flax.typing import Dtype, PrecisionLike
 from typing import Dict, Callable, Sequence, Any, Union
 import einops
 
@@ -18,8 +20,9 @@ class WeightStandardizedConv(nn.Module):
     kernel_size: Sequence[int] = 3
     strides: Union[None, int, Sequence[int]] = 1
     padding: Any = 1
-    dtype: Any = jnp.float32
-    param_dtype: Any = jnp.float32
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    param_dtype: Optional[Dtype] = None
 
     @nn.compact
     def __call__(self, x):
@@ -120,8 +123,8 @@ class SeparableConv(nn.Module):
     use_bias:bool=False
     kernel_init:Callable=kernel_init(1.0)
     padding:str="SAME"
-    dtype: Any = jnp.bfloat16
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
 
     @nn.compact
     def __call__(self, x):
@@ -149,8 +152,8 @@ class ConvLayer(nn.Module):
     kernel_size:tuple=(3, 3)
     strides:tuple=(1, 1)
     kernel_init:Callable=kernel_init(1.0)
-    dtype: Any = jnp.bfloat16
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
 
     def setup(self):
         # conv_type can be "conv", "separable", "conv_transpose"
@@ -199,8 +202,8 @@ class Upsample(nn.Module):
     features:int
     scale:int
     activation:Callable=jax.nn.swish
-    dtype: Any = jnp.bfloat16
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -224,8 +227,8 @@ class Downsample(nn.Module):
     features:int
     scale:int
     activation:Callable=jax.nn.swish
-    dtype: Any = jnp.bfloat16
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -248,3 +251,79 @@ def l2norm(t, axis=1, eps=1e-12):
     denom = jnp.clip(jnp.linalg.norm(t, ord=2, axis=axis, keepdims=True), eps)
     out = t/denom
     return (out)
+
+
+class ResidualBlock(nn.Module):
+    conv_type:str
+    features:int
+    kernel_size:tuple=(3, 3)
+    strides:tuple=(1, 1)
+    padding:str="SAME"
+    activation:Callable=jax.nn.swish
+    direction:str=None
+    res:int=2
+    norm_groups:int=8
+    kernel_init:Callable=kernel_init(1.0)
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+
+    @nn.compact
+    def __call__(self, x:jax.Array, temb:jax.Array, textemb:jax.Array=None, extra_features:jax.Array=None):
+        residual = x
+        # out = nn.GroupNorm(self.norm_groups)(x)
+        out = nn.RMSNorm()(x)
+        out = self.activation(out)
+
+        out = ConvLayer(
+            self.conv_type,
+            features=self.features,
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            kernel_init=self.kernel_init,
+            name="conv1",
+            dtype=self.dtype,
+            precision=self.precision
+        )(out)
+
+        temb = nn.DenseGeneral(
+            features=self.features, 
+            name="temb_projection",
+            dtype=self.dtype,
+            precision=self.precision)(temb)
+        temb = jnp.expand_dims(jnp.expand_dims(temb, 1), 1)
+        # scale, shift = jnp.split(temb, 2, axis=-1)
+        # out = out * (1 + scale) + shift
+        out = out + temb
+
+        # out = nn.GroupNorm(self.norm_groups)(out)
+        out = nn.RMSNorm()(out)
+        out = self.activation(out)
+
+        out = ConvLayer(
+            self.conv_type,
+            features=self.features,
+            kernel_size=self.kernel_size,
+            strides=self.strides,
+            kernel_init=self.kernel_init,
+            name="conv2",
+            dtype=self.dtype,
+            precision=self.precision
+        )(out)
+
+        if residual.shape != out.shape:
+            residual = ConvLayer(
+                self.conv_type,
+                features=self.features,
+                kernel_size=(1, 1),
+                strides=1,
+                kernel_init=self.kernel_init,
+                name="residual_conv",
+                dtype=self.dtype,
+                precision=self.precision
+            )(residual)
+        out = out + residual
+
+        out = jnp.concatenate([out, extra_features], axis=-1) if extra_features is not None else out
+
+        return out
+    

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff/models/simple_unet.py

@@ -1,83 +1,12 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from typing import Dict, Callable, Sequence, Any, Union
+from flax.typing import Dtype, PrecisionLike
+from typing import Dict, Callable, Sequence, Any, Union, Optional
 import einops
 from .common import kernel_init, ConvLayer, Downsample, Upsample, FourierEmbedding, TimeProjection
 from .attention import TransformerBlock
 
-class ResidualBlock(nn.Module):
-    conv_type:str
-    features:int
-    kernel_size:tuple=(3, 3)
-    strides:tuple=(1, 1)
-    padding:str="SAME"
-    activation:Callable=jax.nn.swish
-    direction:str=None
-    res:int=2
-    norm_groups:int=8
-    kernel_init:Callable=kernel_init(1.0)
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGHEST
-
-    @nn.compact
-    def __call__(self, x:jax.Array, temb:jax.Array, textemb:jax.Array=None, extra_features:jax.Array=None):
-        residual = x
-        out = nn.GroupNorm(self.norm_groups)(x)
-        out = self.activation(out)
-
-        out = ConvLayer(
-            self.conv_type,
-            features=self.features,
-            kernel_size=self.kernel_size,
-            strides=self.strides,
-            kernel_init=self.kernel_init,
-            name="conv1",
-            dtype=self.dtype,
-            precision=self.precision
-        )(out)
-
-        temb = nn.DenseGeneral(
-            features=self.features, 
-            name="temb_projection",
-            dtype=self.dtype,
-            precision=self.precision)(temb)
-        temb = jnp.expand_dims(jnp.expand_dims(temb, 1), 1)
-        # scale, shift = jnp.split(temb, 2, axis=-1)
-        # out = out * (1 + scale) + shift
-        out = out + temb
-
-        out = nn.GroupNorm(self.norm_groups)(out)
-        out = self.activation(out)
-
-        out = ConvLayer(
-            self.conv_type,
-            features=self.features,
-            kernel_size=self.kernel_size,
-            strides=self.strides,
-            kernel_init=self.kernel_init,
-            name="conv2",
-            dtype=self.dtype,
-            precision=self.precision
-        )(out)
-
-        if residual.shape != out.shape:
-            residual = ConvLayer(
-                self.conv_type,
-                features=self.features,
-                kernel_size=(1, 1),
-                strides=1,
-                kernel_init=self.kernel_init,
-                name="residual_conv",
-                dtype=self.dtype,
-                precision=self.precision
-            )(residual)
-        out = out + residual
-
-        out = jnp.concatenate([out, extra_features], axis=-1) if extra_features is not None else out
-
-        return out
-    
 class Unet(nn.Module):
     output_channels:int=3
     emb_features:int=64*4,
@@ -87,8 +16,8 @@ class Unet(nn.Module):
     num_middle_res_blocks:int=1,
     activation:Callable = jax.nn.swish
     norm_groups:int=8
-    dtype: Any = jnp.bfloat16
-    precision: Any = jax.lax.Precision.HIGH
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
 
     @nn.compact
     def __call__(self, x, temb, textcontext):

flaxdiff-0.1.6/flaxdiff/trainer/__init__.py

@@ -0,0 +1,2 @@
+from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+from .diffusion_trainer import DiffusionTrainer, TrainState

flaxdiff-0.1.6/flaxdiff/trainer/autoencoder_trainer.py

@@ -0,0 +1,182 @@
+from flax import linen as nn
+import jax
+from typing import Callable
+from dataclasses import field
+import jax.numpy as jnp
+import optax
+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 ..schedulers import NoiseScheduler
+from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+
+from flaxdiff.utils import RandomMarkovState
+
+from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+
+from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
+
+class AutoEncoderTrainer(SimpleTrainer):
+    def __init__(self,
+                 model: nn.Module,
+                 input_shape: Union[int, int, int],
+                 latent_dim: int,
+                 spatial_scale: int,
+                 optimizer: optax.GradientTransformation,
+                 rngs: jax.random.PRNGKey,
+                 name: str = "Autoencoder",
+                 **kwargs
+                 ):
+        super().__init__(
+            model=model,
+            input_shapes={"image": input_shape},
+            optimizer=optimizer,
+            rngs=rngs,
+            name=name,
+            **kwargs
+        )
+        self.latent_dim = latent_dim
+        self.spatial_scale = spatial_scale
+        
+
+    def generate_states(
+        self,
+        optimizer: optax.GradientTransformation,
+        rngs: jax.random.PRNGKey,
+        existing_state: dict = None,
+        existing_best_state: dict = None,
+        model: nn.Module = None,
+        param_transforms: Callable = None
+    ) -> Tuple[TrainState, TrainState]:
+        print("Generating states for DiffusionTrainer")
+        rngs, subkey = jax.random.split(rngs)
+
+        if existing_state == None:
+            input_vars = self.get_input_ones()
+            params = model.init(subkey, **input_vars)
+            new_state = {"params": params, "ema_params": params}
+        else:
+            new_state = existing_state
+
+        if param_transforms is not None:
+            params = param_transforms(params)
+
+        state = TrainState.create(
+            apply_fn=model.apply,
+            params=new_state['params'],
+            ema_params=new_state['ema_params'],
+            tx=optimizer,
+            rngs=rngs,
+            metrics=Metrics.empty()
+        )
+            
+        if existing_best_state is not None:
+            best_state = state.replace(
+                params=existing_best_state['params'], ema_params=existing_best_state['ema_params'])
+        else:
+            best_state = state
+
+        return state, best_state
+
+    def _define_train_step(self, batch_size, null_labels_seq, text_embedder):
+        noise_schedule: NoiseScheduler = self.noise_schedule
+        model = self.model
+        model_output_transform = self.model_output_transform
+        loss_fn = self.loss_fn
+        unconditional_prob = self.unconditional_prob
+
+        # Determine the number of unconditional samples
+        num_unconditional = int(batch_size * unconditional_prob)
+
+        nS, nC = null_labels_seq.shape
+        null_labels_seq = jnp.broadcast_to(
+            null_labels_seq, (batch_size, nS, nC))
+
+        distributed_training = self.distributed_training
+        
+        autoencoder = self.autoencoder
+
+        # @jax.jit
+        def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index):
+            """Train for a single step."""
+            rng_state, subkey = rng_state.get_random_key()
+            subkey = jax.random.fold_in(subkey, local_device_index.reshape())
+            local_rng_state = RandomMarkovState(subkey)
+            
+            images = batch['image']
+            
+            if autoencoder is not None:
+                # Convert the images to latent space
+                local_rng_state, rngs = local_rng_state.get_random_key()
+                images = autoencoder.encode(images, rngs)
+            else:
+                # normalize image
+                images = (images - 127.5) / 127.5
+
+            output = text_embedder(
+                input_ids=batch['input_ids'], attention_mask=batch['attention_mask'])
+            label_seq = output.last_hidden_state
+
+            # Generate random probabilities to decide how much of this batch will be unconditional
+
+            label_seq = jnp.concat(
+                [null_labels_seq[:num_unconditional], label_seq[num_unconditional:]], axis=0)
+
+            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)
+            
+            rates = noise_schedule.get_rates(noise_level)
+            noisy_images, c_in, expected_output = model_output_transform.forward_diffusion(
+                images, noise, rates)
+
+            def model_loss(params):
+                preds = model.apply(params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq)
+                preds = model_output_transform.pred_transform(
+                    noisy_images, preds, rates)
+                nloss = loss_fn(preds, expected_output)
+                # nloss = jnp.mean(nloss, axis=1)
+                nloss *= noise_schedule.get_weights(noise_level)
+                nloss = jnp.mean(nloss)
+                loss = nloss
+                return loss
+            
+            loss, grads = jax.value_and_grad(model_loss)(train_state.params)
+            if distributed_training:
+                grads = jax.lax.pmean(grads, "data")
+                loss = jax.lax.pmean(loss, "data")
+            train_state = train_state.apply_gradients(grads=grads)
+            train_state = train_state.apply_ema(self.ema_decay)
+            return train_state, loss, rng_state
+
+        if distributed_training:
+            train_step = shard_map(train_step, mesh=self.mesh, in_specs=(P(), P(), P('data'), P('data')), 
+                                   out_specs=(P(), P(), P()))
+            train_step = jax.jit(train_step)
+            
+        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']
+        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'
+

flaxdiff-0.1.5/flaxdiff/trainer/__init__.py → flaxdiff-0.1.6/flaxdiff/trainer/diffusion_trainer.py

@@ -15,6 +15,8 @@ from flaxdiff.utils import RandomMarkovState
 
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
 
+from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
+
 class TrainState(SimpleTrainState):
     rngs: jax.random.PRNGKey
     ema_params: dict
@@ -41,6 +43,7 @@ class DiffusionTrainer(SimpleTrainer):
                  unconditional_prob: float = 0.2,
                  name: str = "Diffusion",
                  model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
+                 autoencoder: AutoEncoder = None,
                  **kwargs
                  ):
         super().__init__(
@@ -54,6 +57,8 @@ class DiffusionTrainer(SimpleTrainer):
         self.noise_schedule = noise_schedule
         self.model_output_transform = model_output_transform
         self.unconditional_prob = unconditional_prob
+        
+        self.autoencoder = autoencoder
 
     def generate_states(
         self,
@@ -109,6 +114,8 @@ class DiffusionTrainer(SimpleTrainer):
             null_labels_seq, (batch_size, nS, nC))
 
         distributed_training = self.distributed_training
+        
+        autoencoder = self.autoencoder
 
         # @jax.jit
         def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, local_device_index):
@@ -118,8 +125,14 @@ class DiffusionTrainer(SimpleTrainer):
             local_rng_state = RandomMarkovState(subkey)
             
             images = batch['image']
-            # normalize image
-            images = (images - 127.5) / 127.5
+            
+            if autoencoder is not None:
+                # Convert the images to latent space
+                local_rng_state, rngs = local_rng_state.get_random_key()
+                images = autoencoder.encode(images, rngs)
+            else:
+                # normalize image
+                images = (images - 127.5) / 127.5
 
             output = text_embedder(
                 input_ids=batch['input_ids'], attention_mask=batch['attention_mask'])
@@ -140,8 +153,7 @@ class DiffusionTrainer(SimpleTrainer):
                 images, noise, rates)
 
             def model_loss(params):
-                preds = model.apply(
-                    params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq)
+                preds = model.apply(params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq)
                 preds = model_output_transform.pred_transform(
                     noisy_images, preds, rates)
                 nloss = loss_fn(preds, expected_output)
@@ -182,3 +194,9 @@ class DiffusionTrainer(SimpleTrainer):
         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'
+

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.5
+Version: 0.1.6
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com
@@ -13,6 +13,8 @@ Requires-Dist: clu
 
 # ![](images/logo.jpeg "FlaxDiff")
 
+**This project is partially supported by [Google TPU Research Cloud](https://sites.research.google/trc/about/). I would like to thank the Google Cloud TPU team for providing me with the resources to train the bigger text-conditional models in multi-host distributed settings.**
+
 ## A Versatile and simple Diffusion Library
 
 In recent years, diffusion and score-based multi-step models have revolutionized the generative AI domain. However, the latest research in this field has become highly math-intensive, making it challenging to understand how state-of-the-art diffusion models work and generate such impressive images. Replicating this research in code can be daunting.

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/flaxdiff.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ flaxdiff/models/simple_vit.py
 flaxdiff/models/autoencoder/__init__.py
 flaxdiff/models/autoencoder/autoencoder.py
 flaxdiff/models/autoencoder/diffusers.py
+flaxdiff/models/autoencoder/simple_autoenc.py
 flaxdiff/predictors/__init__.py
 flaxdiff/samplers/__init__.py
 flaxdiff/samplers/common.py
@@ -35,4 +36,6 @@ flaxdiff/schedulers/karras.py
 flaxdiff/schedulers/linear.py
 flaxdiff/schedulers/sqrt.py
 flaxdiff/trainer/__init__.py
+flaxdiff/trainer/autoencoder_trainer.py
+flaxdiff/trainer/diffusion_trainer.py
 flaxdiff/trainer/simple_trainer.py

{flaxdiff-0.1.5 → flaxdiff-0.1.6}/setup.py

@@ -11,7 +11,7 @@ required_packages=[
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.5',
+    version='0.1.6',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',