flaxdiff 0.1.1__py3-none-any.whl

flaxdiff/models/simple_unet.py

@@ -0,0 +1,519 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Dict, Callable, Sequence, Any, Union
+import einops
+from .common import kernel_init
+from .attention import TransformerBlock
+class WeightStandardizedConv(nn.Module):
+    """
+    apply weight standardization  https://arxiv.org/abs/1903.10520
+    """
+    features: int
+    kernel_size: Sequence[int] = 3
+    strides: Union[None, int, Sequence[int]] = 1
+    padding: Any = 1
+    dtype: Any = jnp.float32
+    param_dtype: Any = jnp.float32
+
+    @nn.compact
+    def __call__(self, x):
+        """
+        Applies a weight standardized convolution to the inputs.
+
+        Args:
+          inputs: input data with dimensions (batch, spatial_dims..., features).
+
+        Returns:
+          The convolved data.
+        """
+        x = x.astype(self.dtype)
+
+        conv = nn.Conv(
+            features=self.features,
+            kernel_size=self.kernel_size,
+            strides = self.strides,
+            padding=self.padding,
+            dtype=self.dtype,
+            param_dtype = self.param_dtype,
+            parent=None)
+
+        kernel_init = lambda  rng, x: conv.init(rng,x)['params']['kernel']
+        bias_init = lambda  rng, x: conv.init(rng,x)['params']['bias']
+
+        # standardize kernel
+        kernel = self.param('kernel', kernel_init, x)
+        eps = 1e-5 if self.dtype == jnp.float32 else 1e-3
+        # reduce over dim_out
+        redux = tuple(range(kernel.ndim - 1))
+        mean = jnp.mean(kernel, axis=redux, dtype=self.dtype, keepdims=True)
+        var = jnp.var(kernel, axis=redux, dtype=self.dtype, keepdims=True)
+        standardized_kernel = (kernel - mean)/jnp.sqrt(var + eps)
+
+        bias = self.param('bias',bias_init, x)
+
+        return(conv.apply({'params': {'kernel': standardized_kernel, 'bias': bias}},x))
+
+class PixelShuffle(nn.Module):
+    scale: int
+
+    @nn.compact
+    def __call__(self, x):
+        up = einops.rearrange(
+            x,
+            pattern="b h w (h2 w2 c) -> b (h h2) (w w2) c",
+            h2=self.scale,
+            w2=self.scale,
+        )
+        return up
+
+class TimeEmbedding(nn.Module):
+    features:int
+    nax_positions:int=10000
+
+    def setup(self):
+        half_dim = self.features // 2
+        emb = jnp.log(self.nax_positions) / (half_dim - 1)
+        emb = jnp.exp(-emb * jnp.arange(half_dim, dtype=jnp.float32))
+        self.embeddings = emb
+
+    def __call__(self, x):
+        x = jax.lax.convert_element_type(x, jnp.float32)
+        emb = x[:, None] * self.embeddings[None, :]
+        emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1)
+        return emb
+
+class FourierEmbedding(nn.Module):
+    features:int
+    scale:int = 16
+
+    def setup(self):
+        self.freqs = jax.random.normal(jax.random.PRNGKey(42), (self.features // 2, ), dtype=jnp.float32) * self.scale
+
+    def __call__(self, x):
+        x = jax.lax.convert_element_type(x, jnp.float32)
+        emb = x[:, None] * (2 * jnp.pi * self.freqs)[None, :]
+        emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1)
+        return emb
+
+class TimeProjection(nn.Module):
+    features:int
+    activation:Callable=jax.nn.gelu
+
+    @nn.compact
+    def __call__(self, x):
+        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)
+        x = self.activation(x)
+        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)
+        x = self.activation(x)
+        return x
+
+class SeparableConv(nn.Module):
+    features:int
+    kernel_size:tuple=(3, 3)
+    strides:tuple=(1, 1)
+    use_bias:bool=False
+    kernel_init:Callable=kernel_init(1.0)
+    padding:str="SAME"
+    dtype: Any = jnp.bfloat16
+    precision: Any = jax.lax.Precision.HIGH
+
+    @nn.compact
+    def __call__(self, x):
+        in_features = x.shape[-1]
+        depthwise = nn.Conv(
+            features=in_features, kernel_size=self.kernel_size,
+            strides=self.strides, kernel_init=self.kernel_init,
+            feature_group_count=in_features, use_bias=self.use_bias,
+            padding=self.padding,
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        pointwise = nn.Conv(
+            features=self.features, kernel_size=(1, 1),
+            strides=(1, 1), kernel_init=self.kernel_init,
+            use_bias=self.use_bias,
+            dtype=self.dtype,
+            precision=self.precision
+        )(depthwise)
+        return pointwise
+
+class ConvLayer(nn.Module):
+    conv_type:str
+    features:int
+    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
+
+    def setup(self):
+        # conv_type can be "conv", "separable", "conv_transpose"
+        if self.conv_type == "conv":
+            self.conv = nn.Conv(
+                features=self.features,
+                kernel_size=self.kernel_size,
+                strides=self.strides,
+                kernel_init=self.kernel_init,
+                dtype=self.dtype,
+                precision=self.precision
+            )
+        elif self.conv_type == "w_conv":
+            self.conv = WeightStandardizedConv(
+                features=self.features,
+                kernel_size=self.kernel_size,
+                strides=self.strides,
+                padding="SAME",
+                param_dtype=self.dtype,
+                dtype=self.dtype,
+                precision=self.precision
+            )
+        elif self.conv_type == "separable":
+            self.conv = SeparableConv(
+                features=self.features,
+                kernel_size=self.kernel_size,
+                strides=self.strides,
+                kernel_init=self.kernel_init,
+                dtype=self.dtype,
+                precision=self.precision
+            )
+        elif self.conv_type == "conv_transpose":
+            self.conv = nn.ConvTranspose(
+                features=self.features,
+                kernel_size=self.kernel_size,
+                strides=self.strides,
+                kernel_init=self.kernel_init,
+                dtype=self.dtype,
+                precision=self.precision
+            )
+
+    def __call__(self, x):
+        return self.conv(x)
+
+class Upsample(nn.Module):
+    features:int
+    scale:int
+    activation:Callable=jax.nn.swish
+    dtype: Any = jnp.bfloat16
+    precision: Any = jax.lax.Precision.HIGH
+
+    @nn.compact
+    def __call__(self, x, residual=None):
+        out = x
+        # out = PixelShuffle(scale=self.scale)(out)
+        B, H, W, C = x.shape
+        out = jax.image.resize(x, (B, H * self.scale, W * self.scale, C), method="nearest")
+        out = ConvLayer(
+            "conv",
+            features=self.features,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            dtype=self.dtype,
+            precision=self.precision
+        )(out)
+        if residual is not None:
+            out = jnp.concatenate([out, residual], axis=-1)
+        return out
+
+class Downsample(nn.Module):
+    features:int
+    scale:int
+    activation:Callable=jax.nn.swish
+    dtype: Any = jnp.bfloat16
+    precision: Any = jax.lax.Precision.HIGH
+
+    @nn.compact
+    def __call__(self, x, residual=None):
+        out = ConvLayer(
+            "conv",
+            features=self.features,
+            kernel_size=(3, 3),
+            strides=(2, 2),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        if residual is not None:
+            if residual.shape[1] > out.shape[1]:
+                residual = nn.avg_pool(residual, window_shape=(2, 2), strides=(2, 2), padding="SAME")
+            out = jnp.concatenate([out, residual], axis=-1)
+        return out
+
+
+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: 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):
+    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: Any = jnp.bfloat16
+    precision: Any = jax.lax.Precision.HIGH
+
+    @nn.compact
+    def __call__(self, x, temb, textcontext=None):
+        # 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=kernel_init(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=kernel_init(1.0),
+                    kernel_size=(3, 3),
+                    strides=(1, 1),
+                    activation=self.activation,
+                    norm_groups=self.norm_groups,
+                    dtype=self.dtype,
+                    precision=self.precision
+                )(x, temb)
+                if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
+                    B, H, W, _ = x.shape
+                    if H > TS:
+                        padded_context = jnp.pad(textcontext, ((0, 0), (0, H - TS), (0, 0)), mode='constant', constant_values=0).reshape((B, 1, H, TC))
+                    else:
+                        padded_context = None
+                    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", True),
+                                       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),
+                                       name=f"down_{i}_attention_{j}")(x, padded_context)
+                # 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=kernel_init(1.0),
+                kernel_size=(3, 3),
+                strides=(1, 1),
+                activation=self.activation,
+                norm_groups=self.norm_groups,
+                dtype=self.dtype,
+                precision=self.precision
+            )(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", True),
+                                    use_linear_attention=False,
+                                    use_projection=middle_attention.get("use_projection", False),
+                                    use_self_and_cross=False,
+                                    precision=attention_config.get("precision", self.precision),
+                                    name=f"middle_attention_{j}")(x)
+            x = ResidualBlock(
+                middle_conv_type,
+                name=f"middle_res2_{j}",
+                features=middle_dim_out,
+                kernel_init=kernel_init(1.0),
+                kernel_size=(3, 3),
+                strides=(1, 1),
+                activation=self.activation,
+                norm_groups=self.norm_groups,
+                dtype=self.dtype,
+                precision=self.precision
+            )(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=kernel_init(1.0),
+                    kernel_size=kernel_size,
+                    strides=(1, 1),
+                    activation=self.activation,
+                    norm_groups=self.norm_groups,
+                    dtype=self.dtype,
+                    precision=self.precision
+                )(x, temb)
+                if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
+                    B, H, W, _ = x.shape
+                    if H > TS:
+                        padded_context = jnp.pad(textcontext, ((0, 0), (0, H - TS), (0, 0)), mode='constant', constant_values=0).reshape((B, 1, H, TC))
+                    else:
+                        padded_context = None
+                    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", True),
+                                       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),
+                                       name=f"up_{i}_attention_{j}")(x, padded_context)
+            # 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 = nn.GroupNorm(8)(x)
+        x = ConvLayer(
+            conv_type,
+            features=self.feature_depths[0],
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            kernel_init=kernel_init(0.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=kernel_init(1.0),
+            kernel_size=(3,3),
+            strides=(1, 1),
+            activation=self.activation,
+            norm_groups=self.norm_groups,
+            dtype=self.dtype,
+            precision=self.precision
+        )(x, temb)
+
+        x = nn.GroupNorm(self.norm_groups)(x)
+        x = self.activation(x)
+
+        noise_out = ConvLayer(
+            conv_type,
+            features=3,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            # activation=jax.nn.mish
+            kernel_init=kernel_init(0.0),
+            dtype=self.dtype,
+            precision=self.precision
+        )(x)
+        return noise_out#, attentions

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