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

flaxdiff/models/simple_unet.py

@@ -1,233 +0,0 @@
-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

@@ -1,180 +0,0 @@
-# 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

@@ -1,96 +0,0 @@
-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

@@ -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/samplers/common.py

@@ -1,113 +0,0 @@
-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()):
-        self.model = model
-        self.noise_schedule = noise_schedule
-        self.params = params
-        self.model_output_transform = model_output_transform
-
-        @jax.jit
-        def sample_model(x_t, t):
-            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))
-            x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
-            return x_0, eps, model_output
-        
-        self.sample_model = sample_model
-
-    # Used to sample from the diffusion model
-    def sample_step(self, current_samples:jnp.ndarray, current_step, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
-        # First clip the noisy images
-        # pred_images = clip_images(pred_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)
-        # plotImages(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)
-        return new_samples, state
-
-    def take_next_step(self, current_samples, reconstructed_samples, 
-                 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, image_size=64):
-        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) 
-        return jax.random.normal(rngs, (num_images, image_size, image_size, 3)) * 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))) -> 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")
-            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, 
-                                              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)
-        samples = clip_images(samples)
-        return samples

flaxdiff/samplers/ddim.py

@@ -1,10 +0,0 @@
-import jax.numpy as jnp
-from .common import DiffusionSampler
-from ..utils import MarkovState
-
-class DDIMSampler(DiffusionSampler):
-    def take_next_step(self, 
-                 current_samples, reconstructed_samples, 
-                 pred_noise, current_step, state:MarkovState, next_step=None) -> tuple[jnp.ndarray, MarkovState]:
-        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