flaxdiff 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

flaxdiff/models/simple_unet.py

@@ -1,323 +1,14 @@
 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
+from .common import kernel_init, ConvLayer, Downsample, Upsample, FourierEmbedding, TimeProjection
 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):
+    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}],
@@ -325,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):
@@ -373,12 +64,13 @@ class Unet(nn.Module):
                 )(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", 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, textcontext)
+                                        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),
+                                        only_pure_attention=True,
+                                        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:
@@ -416,6 +108,7 @@ class Unet(nn.Module):
                                     use_projection=middle_attention.get("use_projection", False),
                                     use_self_and_cross=False,
                                     precision=attention_config.get("precision", self.precision),
+                                    only_pure_attention=True,
                                     name=f"middle_attention_{j}")(x, textcontext)
             x = ResidualBlock(
                 middle_conv_type,
@@ -452,12 +145,13 @@ class Unet(nn.Module):
                 )(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", True),
-                                       use_projection=attention_config.get("use_projection", False),
-                                       use_self_and_cross=attention_config.get("use_self_and_cross", True),
+                                        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, textcontext)
+                                        only_pure_attention=True,
+                                        name=f"up_{i}_attention_{j}")(x, textcontext)
             # print("Upscaling ", i, x.shape)
             if i != len(feature_depths) - 1:
                 x = Upsample(
@@ -500,7 +194,7 @@ class Unet(nn.Module):
 
         noise_out = ConvLayer(
             conv_type,
-            features=3,
+            features=self.output_channels,
             kernel_size=(3, 3),
             strides=(1, 1),
             # activation=jax.nn.mish

flaxdiff/trainer/__init__.py

@@ -1,201 +1,2 @@
-import orbax.checkpoint
-import tqdm
-from flax import linen as nn
-import jax
-from typing import Callable
-from dataclasses import field
-import jax.numpy as jnp
-from clu import metrics
-from flax.training import train_state  # Useful dataclass to keep train state
-import optax
-from flax import struct                # Flax dataclasses
-import time
-import os
-import orbax
-from flax.training import orbax_utils
-
-from ..schedulers import NoiseScheduler
-from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
-
-from .simple_trainer import SimpleTrainer, SimpleTrainState
-
-class TrainState(SimpleTrainState):
-    rngs: jax.random.PRNGKey
-    ema_params: dict
-
-    def get_random_key(self):
-        rngs, subkey = jax.random.split(self.rngs)
-        return self.replace(rngs=rngs), subkey
-
-    def apply_ema(self, decay: float = 0.999):
-        new_ema_params = jax.tree_util.tree_map(
-            lambda ema, param: decay * ema + (1 - decay) * param,
-            self.ema_params,
-            self.params,
-        )
-        return self.replace(ema_params=new_ema_params)
-
-class DiffusionTrainer(SimpleTrainer):
-    noise_schedule: NoiseScheduler
-    model_output_transform: DiffusionPredictionTransform
-    ema_decay: float = 0.999
-
-    def __init__(self,
-                 model: nn.Module,
-                 input_shapes: Dict[str, Tuple[int]],
-                 optimizer: optax.GradientTransformation,
-                 noise_schedule: NoiseScheduler,
-                 rngs: jax.random.PRNGKey,
-                 unconditional_prob: float = 0.2,
-                 name: str = "Diffusion",
-                 model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
-                 **kwargs
-                 ):
-        super().__init__(
-            model=model,
-            input_shapes=input_shapes,
-            optimizer=optimizer,
-            rngs=rngs,
-            name=name,
-            **kwargs
-        )
-        self.noise_schedule = noise_schedule
-        self.model_output_transform = model_output_transform
-        self.unconditional_prob = unconditional_prob
-
-    def __init_fn(
-        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]:
-        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 = 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
-
-        def train_step(state: TrainState, batch):
-            """Train for a single step."""
-            images = batch['image']
-            # normalize image
-            images = (images - 127.5) / 127.5
-
-            output = text_embedder(
-                input_ids=batch['input_ids'], attention_mask=batch['attention_mask'])
-            # output = infer(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, state = noise_schedule.generate_timesteps(
-                images.shape[0], state)
-            state, rngs = 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)(state.params)
-            if distributed_training:
-                grads = jax.lax.pmean(grads, "device")
-            state = state.apply_gradients(grads=grads)
-            state = state.apply_ema(self.ema_decay)
-            return state, loss
-        
-        if distributed_training:
-            train_step = jax.pmap(axis_name="device")(train_step)
-        else:
-            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})
-
-
-                        pbar.set_postfix(loss=f'{loss:.4f}')
-                        pbar.update(100)
-            end_time = time.time()
-            self.state = state
-            total_time = end_time - start_time
-            avg_time_per_step = total_time / steps_per_epoch
-            avg_loss = epoch_loss / steps_per_epoch
-            if avg_loss < self.best_loss:
-                self.best_loss = avg_loss
-                self.best_state = state
-                self.save(epoch, best=True)
-            print(f"\n\tEpoch {epoch+1} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}")
-        return self.state
+from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+from .diffusion_trainer import DiffusionTrainer, TrainState