flaxdiff 0.1.36.1__py3-none-any.whl → 0.1.36.3__py3-none-any.whl

flaxdiff/models/simple_unet.py

@@ -0,0 +1,233 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+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, ResidualBlock
+from .attention import TransformerBlock
+from functools import partial
+
+class Unet(nn.Module):
+    output_channels:int=3
+    emb_features:int=64*4,
+    feature_depths:list=[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
+    named_norms: bool = False # This is for backward compatibility reasons; older checkpoints have named norms
+    kernel_init: Callable = partial(kernel_init, dtype=jnp.float32)
+
+    def setup(self):
+        if self.norm_groups > 0:
+            norm = partial(nn.GroupNorm, self.norm_groups)
+            self.conv_out_norm = norm(name="GroupNorm_0") if self.named_norms else norm()
+        else:
+            norm = partial(nn.RMSNorm, 1e-5)
+            self.conv_out_norm = norm()
+        
+    @nn.compact
+    def __call__(self, x, temb, textcontext):
+        # print("embedding features", self.emb_features)
+        temb = FourierEmbedding(features=self.emb_features)(temb)
+        temb = TimeProjection(features=self.emb_features)(temb)
+        
+        _, TS, TC = textcontext.shape
+        
+        # print("time embedding", temb.shape)
+        feature_depths = self.feature_depths
+        attention_configs = self.attention_configs
+
+        conv_type = up_conv_type = down_conv_type = middle_conv_type = "conv"
+        # middle_conv_type = "separable"
+        
+        x = ConvLayer(
+            conv_type,
+            features=self.feature_depths[0],
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            kernel_init=self.kernel_init(scale=1.0),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        downs = [x]
+
+        # Downscaling blocks
+        for i, (dim_out, attention_config) in enumerate(zip(feature_depths, attention_configs)):
+            dim_in = x.shape[-1]
+            # dim_in = dim_out
+            for j in range(self.num_res_blocks):
+                x = ResidualBlock(
+                    down_conv_type,
+                    name=f"down_{i}_residual_{j}",
+                    features=dim_in,
+                    kernel_init=self.kernel_init(scale=1.0),
+                    kernel_size=(3, 3),
+                    strides=(1, 1),
+                    activation=self.activation,
+                    norm_groups=self.norm_groups,
+                    dtype=self.dtype,
+                    precision=self.precision,
+                    named_norms=self.named_norms
+                )(x, temb)
+                if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
+                    x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32),
+                                        dim_head=dim_in // attention_config['heads'],
+                                        use_flash_attention=attention_config.get("flash_attention", False),
+                                        use_projection=attention_config.get("use_projection", False),
+                                        use_self_and_cross=attention_config.get("use_self_and_cross", True),
+                                        precision=attention_config.get("precision", self.precision),
+                                        only_pure_attention=attention_config.get("only_pure_attention", True),
+                                        force_fp32_for_softmax=attention_config.get("force_fp32_for_softmax", False),
+                                        norm_inputs=attention_config.get("norm_inputs", True),
+                                        explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
+                                        kernel_init=self.kernel_init(scale=1.0),
+                                        name=f"down_{i}_attention_{j}")(x, textcontext)
+                # print("down residual for feature level", i, "is of shape", x.shape, "features", dim_in)
+                downs.append(x)
+            if i != len(feature_depths) - 1:
+                # print("Downsample", i, x.shape)
+                x = Downsample(
+                    features=dim_out,
+                    scale=2,
+                    activation=self.activation,
+                    name=f"down_{i}_downsample",
+                    dtype=self.dtype,
+                    precision=self.precision
+                )(x)
+
+        # Middle Blocks
+        middle_dim_out = self.feature_depths[-1]
+        middle_attention = self.attention_configs[-1]
+        for j in range(self.num_middle_res_blocks):
+            x = ResidualBlock(
+                middle_conv_type,
+                name=f"middle_res1_{j}",
+                features=middle_dim_out,
+                kernel_init=self.kernel_init(scale=1.0),
+                kernel_size=(3, 3),
+                strides=(1, 1),
+                activation=self.activation,
+                norm_groups=self.norm_groups,
+                dtype=self.dtype,
+                precision=self.precision,
+                named_norms=self.named_norms
+            )(x, temb)
+            if middle_attention is not None and j == self.num_middle_res_blocks - 1:   # Apply attention only on the last block
+                x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), 
+                                    dim_head=middle_dim_out // middle_attention['heads'],
+                                    use_flash_attention=middle_attention.get("flash_attention", False),
+                                    use_linear_attention=False,
+                                    use_projection=middle_attention.get("use_projection", False),
+                                    use_self_and_cross=False,
+                                    precision=middle_attention.get("precision", self.precision),
+                                    only_pure_attention=middle_attention.get("only_pure_attention", True),
+                                    force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
+                                    norm_inputs=middle_attention.get("norm_inputs", True),
+                                    explicitly_add_residual=middle_attention.get("explicitly_add_residual", True),
+                                    kernel_init=self.kernel_init(scale=1.0),
+                                    name=f"middle_attention_{j}")(x, textcontext)
+            x = ResidualBlock(
+                middle_conv_type,
+                name=f"middle_res2_{j}",
+                features=middle_dim_out,
+                kernel_init=self.kernel_init(scale=1.0),
+                kernel_size=(3, 3),
+                strides=(1, 1),
+                activation=self.activation,
+                norm_groups=self.norm_groups,
+                dtype=self.dtype,
+                precision=self.precision,
+                named_norms=self.named_norms
+            )(x, temb)
+
+        # Upscaling Blocks
+        for i, (dim_out, attention_config) in enumerate(zip(reversed(feature_depths), reversed(attention_configs))):
+            # print("Upscaling", i, "features", dim_out)
+            for j in range(self.num_res_blocks):
+                x = jnp.concatenate([x, downs.pop()], axis=-1)
+                # print("concat==> ", i, "concat", x.shape)
+                # kernel_size = (1 + 2 * (j + 1), 1 + 2 * (j + 1))
+                kernel_size = (3, 3)
+                x = ResidualBlock(
+                    up_conv_type,# if j == 0 else "separable",
+                    name=f"up_{i}_residual_{j}",
+                    features=dim_out,
+                    kernel_init=self.kernel_init(scale=1.0),
+                    kernel_size=kernel_size,
+                    strides=(1, 1),
+                    activation=self.activation,
+                    norm_groups=self.norm_groups,
+                    dtype=self.dtype,
+                    precision=self.precision,
+                    named_norms=self.named_norms
+                )(x, temb)
+                if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
+                    x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), 
+                                        dim_head=dim_out // attention_config['heads'],
+                                        use_flash_attention=attention_config.get("flash_attention", False),
+                                        use_projection=attention_config.get("use_projection", False),
+                                        use_self_and_cross=attention_config.get("use_self_and_cross", True),
+                                        precision=attention_config.get("precision", self.precision),
+                                        only_pure_attention=attention_config.get("only_pure_attention", True),
+                                        force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
+                                        norm_inputs=attention_config.get("norm_inputs", True),
+                                        explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
+                                        kernel_init=self.kernel_init(scale=1.0),
+                                        name=f"up_{i}_attention_{j}")(x, textcontext)
+            # print("Upscaling ", i, x.shape)
+            if i != len(feature_depths) - 1:
+                x = Upsample(
+                    features=feature_depths[-i],
+                    scale=2,
+                    activation=self.activation,
+                    name=f"up_{i}_upsample",
+                    dtype=self.dtype,
+                    precision=self.precision
+                )(x)
+
+        # x = self.last_up_norm(x)
+        x = ConvLayer(
+            conv_type,
+            features=self.feature_depths[0],
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            kernel_init=self.kernel_init(scale=1.0),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+    
+        x = jnp.concatenate([x, downs.pop()], axis=-1)
+
+        x = ResidualBlock(
+            conv_type,
+            name="final_residual",
+            features=self.feature_depths[0],
+            kernel_init=self.kernel_init(scale=1.0),
+            kernel_size=(3,3),
+            strides=(1, 1),
+            activation=self.activation,
+            norm_groups=self.norm_groups,
+            dtype=self.dtype,
+            precision=self.precision,
+            named_norms=self.named_norms
+        )(x, temb)
+
+        x = self.conv_out_norm(x)
+        x = self.activation(x)
+
+        noise_out = ConvLayer(
+            conv_type,
+            features=self.output_channels,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            # activation=jax.nn.mish
+            kernel_init=self.kernel_init(scale=0.0),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        return noise_out#, attentions

flaxdiff/models/simple_vit.py

@@ -0,0 +1,180 @@
+# simple_vit.py
+
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Callable, Any, Optional, Tuple
+from .simple_unet import FourierEmbedding, TimeProjection, ConvLayer, kernel_init
+from .attention import TransformerBlock
+from flaxdiff.models.simple_unet import FourierEmbedding, TimeProjection, ConvLayer, kernel_init, ResidualBlock
+import einops
+from flax.typing import Dtype, PrecisionLike
+from functools import partial
+
+def unpatchify(x, channels=3):
+    patch_size = int((x.shape[2] // channels) ** 0.5)
+    h = w = int(x.shape[1] ** .5)
+    assert h * w == x.shape[1] and patch_size ** 2 * channels == x.shape[2], f"Invalid shape: {x.shape}, should be {h*w}, {patch_size**2*channels}"
+    x = einops.rearrange(x, 'B (h w) (p1 p2 C) -> B (h p1) (w p2) C', h=h, p1=patch_size, p2=patch_size)
+    return x
+
+class PatchEmbedding(nn.Module):
+    patch_size: int
+    embedding_dim: int
+    dtype: Any = jnp.float32
+    precision: Any = jax.lax.Precision.HIGH
+    kernel_init: Callable = partial(kernel_init, 1.0)
+
+    @nn.compact
+    def __call__(self, x):
+        batch, height, width, channels = x.shape
+        assert height % self.patch_size == 0 and width % self.patch_size == 0, "Image dimensions must be divisible by patch size"
+        
+        x = nn.Conv(features=self.embedding_dim, 
+                    kernel_size=(self.patch_size, self.patch_size), 
+                    strides=(self.patch_size, self.patch_size),
+                    dtype=self.dtype,
+                    kernel_init=self.kernel_init(),
+                    precision=self.precision)(x)
+        x = jnp.reshape(x, (batch, -1, self.embedding_dim))
+        return x
+
+class PositionalEncoding(nn.Module):
+    max_len: int
+    embedding_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        pe = self.param('pos_encoding',
+                        jax.nn.initializers.zeros,
+                        (1, self.max_len, self.embedding_dim))
+        return x + pe[:, :x.shape[1], :]
+
+class UViT(nn.Module):
+    output_channels:int=3
+    patch_size: int = 16
+    emb_features:int=768,
+    num_layers: int = 12
+    num_heads: int = 12
+    dropout_rate: float = 0.1
+    dtype: Any = jnp.float32
+    precision: Any = jax.lax.Precision.HIGH
+    use_projection: bool = False
+    use_flash_attention: bool = False
+    use_self_and_cross: bool = False
+    force_fp32_for_softmax: bool = True
+    activation:Callable = jax.nn.swish
+    norm_groups:int=8
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    kernel_init: Callable = partial(kernel_init, scale=1.0)
+    add_residualblock_output: bool = False
+    norm_inputs: bool = False
+    explicitly_add_residual: bool = True
+
+    def setup(self):
+        if self.norm_groups > 0:
+            self.norm = partial(nn.GroupNorm, self.norm_groups)
+        else:
+            self.norm = partial(nn.RMSNorm, 1e-5)
+            
+    @nn.compact
+    def __call__(self, x, temb, textcontext=None):
+        # Time embedding
+        temb = FourierEmbedding(features=self.emb_features)(temb)
+        temb = TimeProjection(features=self.emb_features)(temb)
+        
+        original_img = x
+
+        # Patch embedding
+        x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
+                           dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init)(x)
+        num_patches = x.shape[1]
+        
+        context_emb = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
+                               dtype=self.dtype, precision=self.precision)(textcontext)
+        num_text_tokens = textcontext.shape[1]
+        
+        # print(f'Shape of x after patch embedding: {x.shape}, numPatches: {num_patches}, temb: {temb.shape}, context_emb: {context_emb.shape}')
+        
+        # Add time embedding
+        temb = jnp.expand_dims(temb, axis=1)
+        x = jnp.concatenate([x, temb, context_emb], axis=1)
+        # print(f'Shape of x after time embedding: {x.shape}')
+        
+        # Add positional encoding
+        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.emb_features)(x)
+        
+        # print(f'Shape of x after positional encoding: {x.shape}')
+        
+        skips = []
+        # In blocks
+        for i in range(self.num_layers // 2):
+            x = TransformerBlock(heads=self.num_heads, dim_head=self.emb_features // self.num_heads, 
+                                 dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, 
+                                 use_flash_attention=self.use_flash_attention, use_self_and_cross=False, force_fp32_for_softmax=self.force_fp32_for_softmax, 
+                                 only_pure_attention=False,
+                                 norm_inputs=self.norm_inputs,
+                                 explicitly_add_residual=self.explicitly_add_residual,
+                                 kernel_init=self.kernel_init())(x)
+            skips.append(x)
+            
+        # Middle block
+        x = TransformerBlock(heads=self.num_heads, dim_head=self.emb_features // self.num_heads, 
+                             dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, 
+                             use_flash_attention=self.use_flash_attention, use_self_and_cross=False, force_fp32_for_softmax=self.force_fp32_for_softmax, 
+                             only_pure_attention=False,
+                            norm_inputs=self.norm_inputs,
+                            explicitly_add_residual=self.explicitly_add_residual,
+                             kernel_init=self.kernel_init())(x)
+        
+        # # Out blocks
+        for i in range(self.num_layers // 2):
+            x = jnp.concatenate([x, skips.pop()], axis=-1)
+            x = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
+                                   dtype=self.dtype, precision=self.precision)(x)
+            x = TransformerBlock(heads=self.num_heads, dim_head=self.emb_features // self.num_heads, 
+                                 dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, 
+                                 use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, 
+                                 only_pure_attention=False,
+                                 norm_inputs=self.norm_inputs,
+                                 explicitly_add_residual=self.explicitly_add_residual,
+                                 kernel_init=self.kernel_init())(x)
+        
+        # print(f'Shape of x after transformer blocks: {x.shape}')
+        x = self.norm()(x)
+        
+        patch_dim = self.patch_size ** 2 * self.output_channels
+        x = nn.Dense(features=patch_dim, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
+        x = x[:, 1 + num_text_tokens:, :]
+        x = unpatchify(x, channels=self.output_channels)
+        
+        if self.add_residualblock_output:
+            # Concatenate the original image
+            x = jnp.concatenate([original_img, x], axis=-1)
+            
+            x = ConvLayer(
+                "conv",
+                features=64,
+                kernel_size=(3, 3),
+                strides=(1, 1),
+                # activation=jax.nn.mish
+                kernel_init=self.kernel_init(scale=0.0),
+                dtype=self.dtype,
+                precision=self.precision
+            )(x)
+
+            x = self.norm()(x)
+            x = self.activation(x)
+
+        x = ConvLayer(
+            "conv",
+            features=self.output_channels,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            # activation=jax.nn.mish
+            kernel_init=self.kernel_init(scale=0.0),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        return x

flaxdiff/predictors/__init__.py

@@ -0,0 +1,96 @@
+from typing import Union
+import jax.numpy as jnp
+from ..schedulers import NoiseScheduler, GeneralizedNoiseScheduler
+
+############################################################################################################
+# Prediction Transforms
+############################################################################################################
+
+class DiffusionPredictionTransform():
+    def pred_transform(self, x_t, preds, rates) -> jnp.ndarray:
+        return preds
+    
+    def __call__(self, x_t, preds, current_step, noise_schedule:NoiseScheduler) -> Union[jnp.ndarray, jnp.ndarray]:
+        rates = noise_schedule.get_rates(current_step)
+        preds = self.pred_transform(x_t, preds, rates)
+        x_0, epsilon = self.backward_diffusion(x_t, preds, rates)
+        return x_0, epsilon
+    
+    def forward_diffusion(self, x_0, epsilon, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+        signal_rate, noise_rate = rates
+        x_t = signal_rate * x_0 + noise_rate * epsilon
+        expected_output = self.get_target(x_0, epsilon, (signal_rate, noise_rate))
+        c_in = self.get_input_scale((signal_rate, noise_rate))
+        return x_t, c_in, expected_output
+    
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        raise NotImplementedError
+    
+    def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
+        return x_0
+    
+    def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray:
+        return 1
+
+class EpsilonPredictionTransform(DiffusionPredictionTransform):
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        # preds is the predicted noise
+        epsilon = preds
+        signal_rates, noise_rates = rates
+        x_0 = (x_t - epsilon * noise_rates) / signal_rates
+        return x_0, epsilon
+    
+    def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
+        return epsilon
+
+class DirectPredictionTransform(DiffusionPredictionTransform):
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        # Here the model predicts x_0 directly
+        x_0 = preds
+        signal_rate, noise_rate = rates
+        epsilon = (x_t - x_0 * signal_rate) / noise_rate
+        return x_0, epsilon
+
+class VPredictionTransform(DiffusionPredictionTransform):
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        # here the model output's V = sqrt_alpha_t * epsilon - sqrt_one_minus_alpha_t * x_0
+        # where epsilon is the noise
+        # x_0 is the current sample
+        v = preds
+        signal_rate, noise_rate = rates
+        variance = signal_rate ** 2 + noise_rate ** 2
+        v = v * jnp.sqrt(variance)
+        x_0 = signal_rate * x_t - noise_rate * v
+        eps_0 = signal_rate * v + noise_rate * x_t
+        return x_0 / variance, eps_0 / variance
+    
+    def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
+        signal_rate, noise_rate = rates
+        v = signal_rate * epsilon - noise_rate * x_0
+        variance = signal_rate**2 + noise_rate**2
+        return v / jnp.sqrt(variance)
+    
+class KarrasPredictionTransform(DiffusionPredictionTransform):
+    def __init__(self, sigma_data=0.5) -> None:
+        super().__init__()
+        self.sigma_data = sigma_data
+
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        x_0 = preds
+        signal_rate, noise_rate = rates
+        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:
+        _, 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 = 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:
+        _, sigma = rates
+        c_in = 1 / jnp.sqrt(self.sigma_data ** 2 + sigma ** 2)
+        return c_in

flaxdiff/samplers/__init__.py

@@ -0,0 +1,7 @@
+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/samplers/common.py

@@ -0,0 +1,165 @@
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+import tqdm
+from typing import Union
+from ..schedulers import NoiseScheduler
+from ..utils import RandomMarkovState, MarkovState, clip_images
+from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+
+class DiffusionSampler():
+    model:nn.Module
+    noise_schedule:NoiseScheduler
+    params:dict
+    model_output_transform:DiffusionPredictionTransform
+
+    def __init__(self, model:nn.Module, params:dict,  
+                 noise_schedule:NoiseScheduler, 
+                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform(),
+                 guidance_scale:float = 0.0,
+                 null_labels_seq:jax.Array=None,
+                 autoencoder=None,
+                 image_size=256,
+                 autoenc_scale_reduction=8,
+                 autoenc_latent_channels=4,
+                 ):
+        self.model = model
+        self.noise_schedule = noise_schedule
+        self.params = params
+        self.model_output_transform = model_output_transform
+        self.guidance_scale = guidance_scale
+        self.image_size = image_size
+        self.autoenc_scale_reduction = autoenc_scale_reduction
+        self.autoencoder = autoencoder
+        self.autoenc_latent_channels = autoenc_latent_channels
+        
+        if self.guidance_scale > 0:
+            # 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):
+                # Concatenate unconditional and conditional inputs
+                x_t_cat = jnp.concatenate([x_t] * 2, axis=0)
+                t_cat = jnp.concatenate([t] * 2, axis=0)
+                rates_cat = self.noise_schedule.get_rates(t_cat)
+                c_in_cat = self.model_output_transform.get_input_scale(rates_cat)
+                
+                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)
+                # 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)
+                
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+        else:
+            # Unconditional sampling
+            def sample_model(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)
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+            
+        # if jax.device_count() > 1:
+        #     mesh = jax.sharding.Mesh(jax.devices(), 'data')
+        #     sample_model = shard_map(sample_model, mesh=mesh, in_specs=(P('data'), P('data'), P('data')),
+        #                              out_specs=(P('data'), P('data'), P('data')))
+        sample_model = jax.jit(sample_model)
+        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]:
+        # 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)
+        # 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
+                             )
+        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]:
+        # 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
+    
+    def scale_steps(self, steps):
+        scale_factor = self.noise_schedule.max_timesteps / 1000
+        return steps * scale_factor
+
+    def get_steps(self, start_step, end_step, diffusion_steps):
+        step_range = start_step - end_step
+        if diffusion_steps is None or diffusion_steps == 0:
+            diffusion_steps = start_step - end_step
+        diffusion_steps = min(diffusion_steps, step_range)
+        steps = jnp.linspace(end_step, start_step, diffusion_steps, dtype=jnp.int16)[::-1]
+        return steps
+    
+    def get_initial_samples(self, num_images, rngs:jax.random.PRNGKey, start_step):
+        start_step = self.scale_steps(start_step)
+        alpha_n, sigma_n = self.noise_schedule.get_rates(start_step)
+        variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2) 
+        image_size = self.image_size
+        image_channels = 3
+        if self.autoencoder is not None:
+            image_size = image_size // self.autoenc_scale_reduction
+            image_channels = self.autoenc_latent_channels
+        return jax.random.normal(rngs, (num_images, image_size, image_size, image_channels)) * variance
+
+    def generate_images(self,
+                        num_images=16, 
+                        diffusion_steps=1000, 
+                        start_step:int = None,
+                        end_step:int = 0,
+                        steps_override=None,
+                        priors=None, 
+                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42)),
+                        model_conditioning_inputs:tuple=()
+                        ) -> jnp.ndarray:
+        if priors is None:
+            rngstate, newrngs = rngstate.get_random_key()
+            samples = self.get_initial_samples(num_images, newrngs, start_step)
+        else:
+            print("Using priors")
+            if self.autoencoder is not None:
+                priors = self.autoencoder.encode(priors)
+            samples = priors
+
+        # @jax.jit
+        def sample_step(state:RandomMarkovState, samples, current_step, next_step):
+            samples, state = self.sample_step(current_samples=samples,
+                                              current_step=current_step, 
+                                              model_conditioning_inputs=model_conditioning_inputs,
+                                              state=state, next_step=next_step)
+            return samples, state
+
+        if start_step is None:
+            start_step = self.noise_schedule.max_timesteps
+
+        if steps_override is not None:
+            steps = steps_override
+        else:
+            steps = self.get_steps(start_step, end_step, diffusion_steps)
+
+        # print("Sampling steps", steps)
+        for i in tqdm.tqdm(range(0, len(steps))):
+            current_step = self.scale_steps(steps[i])
+            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)
+            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)
+        if self.autoencoder is not None:
+            samples = self.autoencoder.decode(samples)
+        samples = clip_images(samples)
+        return samples