diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index b3f60c6..bb482f2 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -26,14 +26,14 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip build - pip install flake8 pytest - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + pip install flake8 + pip install -e .[test] - name: Lint with flake8 run: | # stop the build if there are Python syntax errors or undefined names flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics - # - name: Test with pytest - # run: | - # pytest + - name: Test with pytest + run: | + JAX_PLATFORMS=cpu pytest -m "not network" -q diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 70962fd..b67800a 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -8,9 +8,8 @@ name: Upload Python Package +# Publish only on releases - pushes to main must not ship untested code to PyPI on: - push: - branches: [ "main" ] release: types: [published] diff --git a/.gitignore b/.gitignore index d3b2f90..c493820 100644 --- a/.gitignore +++ b/.gitignore @@ -22,4 +22,6 @@ gcsfuse.yml *.arrow artifacts *.parquet -datasets/datasets \ No newline at end of file +datasets/datasets +# local working docs +PLAN.md diff --git a/README.md b/README.md index 0630a24..3eee41c 100644 --- a/README.md +++ b/README.md @@ -108,15 +108,18 @@ Implemented in `flaxdiff.models`: ## Installation -To install FlaxDiff, you need to have Python 3.10 or higher. Install the required dependencies using: +To install FlaxDiff, you need to have Python 3.10 or higher: ```bash -pip install -r requirements.txt +pip install flaxdiff ``` -The models were trained and tested with jax==0.4.28 and flax==0.8.4. However, when I updated to the latest jax==0.4.30 and flax==0.8.5, -the models stopped training. There seems to have been some major change breaking the training dynamics and therefore I would recommend -sticking to the versions mentioned in the requirements.txt +Or for development, clone the repo and install in editable mode with the test dependencies: + +```bash +pip install -e .[test] +pytest -m "not network" +``` ## Getting Started diff --git a/flaxdiff/models/autoencoder/diffusers.py b/flaxdiff/models/autoencoder/diffusers.py index 6a409fd..43ee03c 100644 --- a/flaxdiff/models/autoencoder/diffusers.py +++ b/flaxdiff/models/autoencoder/diffusers.py @@ -2,70 +2,67 @@ import jax.numpy as jnp from flax import linen as nn from .autoencoder import AutoEncoder +from .vae import FlaxEncoder, FlaxDecoder, load_pretrained_vae """ -This module contains an Autoencoder implementation which uses the Stable Diffusion VAE model from the HuggingFace Diffusers library. -The actual model was not trained by me, but was taken from the HuggingFace model hub. -I have only implemented the wrapper around the diffusers pipeline to make it compatible with our library -All credits for the model go to the developers of Stable Diffusion VAE and all credits for the pipeline go to the developers of the Diffusers library. +This module contains an Autoencoder implementation which uses the Stable Diffusion VAE model from the HuggingFace model hub. +The actual model was not trained by me. I have only implemented the wrapper (and vendored the diffusers Flax modules +in vae.py, as diffusers dropped Flax support upstream) to make it compatible with our library. +All credits for the model go to the developers of Stable Diffusion VAE and for the modules to the Diffusers library. """ class StableDiffusionVAE(AutoEncoder): def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16): - - from diffusers.models.vae_flax import FlaxEncoder, FlaxDecoder - from diffusers import FlaxStableDiffusionPipeline, FlaxAutoencoderKL - - vae, params = FlaxAutoencoderKL.from_pretrained( - modelname, - # revision=revision, - dtype=dtype, - ) - + + pretrained = load_pretrained_vae(modelname) + config = pretrained["config"] + params = pretrained["params"] + self.modelname = modelname self.revision = revision self.dtype = dtype - + enc = FlaxEncoder( - in_channels=vae.config.in_channels, - out_channels=vae.config.latent_channels, - down_block_types=vae.config.down_block_types, - block_out_channels=vae.config.block_out_channels, - layers_per_block=vae.config.layers_per_block, - act_fn=vae.config.act_fn, - norm_num_groups=vae.config.norm_num_groups, + in_channels=config["in_channels"], + out_channels=config["latent_channels"], + down_block_types=config["down_block_types"], + block_out_channels=config["block_out_channels"], + layers_per_block=config["layers_per_block"], + act_fn=config["act_fn"], + norm_num_groups=config["norm_num_groups"], double_z=True, - dtype=vae.dtype, + dtype=dtype, ) - + dec = FlaxDecoder( - in_channels=vae.config.latent_channels, - out_channels=vae.config.out_channels, - up_block_types=vae.config.up_block_types, - block_out_channels=vae.config.block_out_channels, - layers_per_block=vae.config.layers_per_block, - norm_num_groups=vae.config.norm_num_groups, - act_fn=vae.config.act_fn, - dtype=vae.dtype, + in_channels=config["latent_channels"], + out_channels=config["out_channels"], + up_block_types=config["up_block_types"], + block_out_channels=config["block_out_channels"], + layers_per_block=config["layers_per_block"], + norm_num_groups=config["norm_num_groups"], + act_fn=config["act_fn"], + dtype=dtype, ) - + quant_conv = nn.Conv( - 2 * vae.config.latent_channels, + 2 * config["latent_channels"], kernel_size=(1, 1), strides=(1, 1), padding="VALID", - dtype=vae.dtype, + dtype=dtype, ) post_quant_conv = nn.Conv( - vae.config.latent_channels, + config["latent_channels"], kernel_size=(1, 1), strides=(1, 1), padding="VALID", - dtype=vae.dtype, + dtype=dtype, ) - - scaling_factor = vae.scaling_factor + + # Older VAE configs predate the scaling_factor key; 0.18215 is the SD default + scaling_factor = config.get("scaling_factor", 0.18215) print(f"Scaling factor: {scaling_factor}") def encode_single_frame(images, rngkey: jax.random.PRNGKey = None): diff --git a/flaxdiff/models/autoencoder/vae.py b/flaxdiff/models/autoencoder/vae.py new file mode 100644 index 0000000..5dbc126 --- /dev/null +++ b/flaxdiff/models/autoencoder/vae.py @@ -0,0 +1,686 @@ +""" +Flax VAE (Stable Diffusion AutoencoderKL) modules, vendored from +huggingface/diffusers v0.29.2 `src/diffusers/models/vae_flax.py` (Apache-2.0). + +Vendored because diffusers removed all Flax/JAX support upstream +(huggingface/diffusers#14169), which would have broken our latent diffusion +path on any fresh install. Only the plain linen modules are kept; the +diffusers ConfigMixin/FlaxModelMixin machinery is replaced by +`load_pretrained_vae` below, which pulls the config and msgpack weights +straight from the HuggingFace Hub. +""" + +import json +import math +from functools import partial +from typing import Tuple + +import flax.linen as nn +import jax +import jax.numpy as jnp + +class FlaxUpsample2D(nn.Module): + """ + Flax implementation of 2D Upsample layer + + Args: + in_channels (`int`): + Input channels + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.conv = nn.Conv( + self.in_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + def __call__(self, hidden_states): + batch, height, width, channels = hidden_states.shape + hidden_states = jax.image.resize( + hidden_states, + shape=(batch, height * 2, width * 2, channels), + method="nearest", + ) + hidden_states = self.conv(hidden_states) + return hidden_states + + +class FlaxDownsample2D(nn.Module): + """ + Flax implementation of 2D Downsample layer + + Args: + in_channels (`int`): + Input channels + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.conv = nn.Conv( + self.in_channels, + kernel_size=(3, 3), + strides=(2, 2), + padding="VALID", + dtype=self.dtype, + ) + + def __call__(self, hidden_states): + pad = ((0, 0), (0, 1), (0, 1), (0, 0)) # pad height and width dim + hidden_states = jnp.pad(hidden_states, pad_width=pad) + hidden_states = self.conv(hidden_states) + return hidden_states + + +class FlaxResnetBlock2D(nn.Module): + """ + Flax implementation of 2D Resnet Block. + + Args: + in_channels (`int`): + Input channels + out_channels (`int`): + Output channels + dropout (:obj:`float`, *optional*, defaults to 0.0): + Dropout rate + groups (:obj:`int`, *optional*, defaults to `32`): + The number of groups to use for group norm. + use_nin_shortcut (:obj:`bool`, *optional*, defaults to `None`): + Whether to use `nin_shortcut`. This activates a new layer inside ResNet block + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + out_channels: int = None + dropout: float = 0.0 + groups: int = 32 + use_nin_shortcut: bool = None + dtype: jnp.dtype = jnp.float32 + + def setup(self): + out_channels = self.in_channels if self.out_channels is None else self.out_channels + + self.norm1 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) + self.conv1 = nn.Conv( + out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + self.norm2 = nn.GroupNorm(num_groups=self.groups, epsilon=1e-6) + self.dropout_layer = nn.Dropout(self.dropout) + self.conv2 = nn.Conv( + out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + use_nin_shortcut = self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut + + self.conv_shortcut = None + if use_nin_shortcut: + self.conv_shortcut = nn.Conv( + out_channels, + kernel_size=(1, 1), + strides=(1, 1), + padding="VALID", + dtype=self.dtype, + ) + + def __call__(self, hidden_states, deterministic=True): + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = nn.swish(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = nn.swish(hidden_states) + hidden_states = self.dropout_layer(hidden_states, deterministic) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + return hidden_states + residual + + +class FlaxAttentionBlock(nn.Module): + r""" + Flax Convolutional based multi-head attention block for diffusion-based VAE. + + Parameters: + channels (:obj:`int`): + Input channels + num_head_channels (:obj:`int`, *optional*, defaults to `None`): + Number of attention heads + num_groups (:obj:`int`, *optional*, defaults to `32`): + The number of groups to use for group norm + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + + """ + + channels: int + num_head_channels: int = None + num_groups: int = 32 + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.num_heads = self.channels // self.num_head_channels if self.num_head_channels is not None else 1 + + dense = partial(nn.Dense, self.channels, dtype=self.dtype) + + self.group_norm = nn.GroupNorm(num_groups=self.num_groups, epsilon=1e-6) + self.query, self.key, self.value = dense(), dense(), dense() + self.proj_attn = dense() + + def transpose_for_scores(self, projection): + new_projection_shape = projection.shape[:-1] + (self.num_heads, -1) + # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) + new_projection = projection.reshape(new_projection_shape) + # (B, T, H, D) -> (B, H, T, D) + new_projection = jnp.transpose(new_projection, (0, 2, 1, 3)) + return new_projection + + def __call__(self, hidden_states): + residual = hidden_states + batch, height, width, channels = hidden_states.shape + + hidden_states = self.group_norm(hidden_states) + + hidden_states = hidden_states.reshape((batch, height * width, channels)) + + query = self.query(hidden_states) + key = self.key(hidden_states) + value = self.value(hidden_states) + + # transpose + query = self.transpose_for_scores(query) + key = self.transpose_for_scores(key) + value = self.transpose_for_scores(value) + + # compute attentions + scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads)) + attn_weights = jnp.einsum("...qc,...kc->...qk", query * scale, key * scale) + attn_weights = nn.softmax(attn_weights, axis=-1) + + # attend to values + hidden_states = jnp.einsum("...kc,...qk->...qc", value, attn_weights) + + hidden_states = jnp.transpose(hidden_states, (0, 2, 1, 3)) + new_hidden_states_shape = hidden_states.shape[:-2] + (self.channels,) + hidden_states = hidden_states.reshape(new_hidden_states_shape) + + hidden_states = self.proj_attn(hidden_states) + hidden_states = hidden_states.reshape((batch, height, width, channels)) + hidden_states = hidden_states + residual + return hidden_states + + +class FlaxDownEncoderBlock2D(nn.Module): + r""" + Flax Resnet blocks-based Encoder block for diffusion-based VAE. + + Parameters: + in_channels (:obj:`int`): + Input channels + out_channels (:obj:`int`): + Output channels + dropout (:obj:`float`, *optional*, defaults to 0.0): + Dropout rate + num_layers (:obj:`int`, *optional*, defaults to 1): + Number of Resnet layer block + resnet_groups (:obj:`int`, *optional*, defaults to `32`): + The number of groups to use for the Resnet block group norm + add_downsample (:obj:`bool`, *optional*, defaults to `True`): + Whether to add downsample layer + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + out_channels: int + dropout: float = 0.0 + num_layers: int = 1 + resnet_groups: int = 32 + add_downsample: bool = True + dtype: jnp.dtype = jnp.float32 + + def setup(self): + resnets = [] + for i in range(self.num_layers): + in_channels = self.in_channels if i == 0 else self.out_channels + + res_block = FlaxResnetBlock2D( + in_channels=in_channels, + out_channels=self.out_channels, + dropout=self.dropout, + groups=self.resnet_groups, + dtype=self.dtype, + ) + resnets.append(res_block) + self.resnets = resnets + + if self.add_downsample: + self.downsamplers_0 = FlaxDownsample2D(self.out_channels, dtype=self.dtype) + + def __call__(self, hidden_states, deterministic=True): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, deterministic=deterministic) + + if self.add_downsample: + hidden_states = self.downsamplers_0(hidden_states) + + return hidden_states + + +class FlaxUpDecoderBlock2D(nn.Module): + r""" + Flax Resnet blocks-based Decoder block for diffusion-based VAE. + + Parameters: + in_channels (:obj:`int`): + Input channels + out_channels (:obj:`int`): + Output channels + dropout (:obj:`float`, *optional*, defaults to 0.0): + Dropout rate + num_layers (:obj:`int`, *optional*, defaults to 1): + Number of Resnet layer block + resnet_groups (:obj:`int`, *optional*, defaults to `32`): + The number of groups to use for the Resnet block group norm + add_upsample (:obj:`bool`, *optional*, defaults to `True`): + Whether to add upsample layer + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + out_channels: int + dropout: float = 0.0 + num_layers: int = 1 + resnet_groups: int = 32 + add_upsample: bool = True + dtype: jnp.dtype = jnp.float32 + + def setup(self): + resnets = [] + for i in range(self.num_layers): + in_channels = self.in_channels if i == 0 else self.out_channels + res_block = FlaxResnetBlock2D( + in_channels=in_channels, + out_channels=self.out_channels, + dropout=self.dropout, + groups=self.resnet_groups, + dtype=self.dtype, + ) + resnets.append(res_block) + + self.resnets = resnets + + if self.add_upsample: + self.upsamplers_0 = FlaxUpsample2D(self.out_channels, dtype=self.dtype) + + def __call__(self, hidden_states, deterministic=True): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, deterministic=deterministic) + + if self.add_upsample: + hidden_states = self.upsamplers_0(hidden_states) + + return hidden_states + + +class FlaxUNetMidBlock2D(nn.Module): + r""" + Flax Unet Mid-Block module. + + Parameters: + in_channels (:obj:`int`): + Input channels + dropout (:obj:`float`, *optional*, defaults to 0.0): + Dropout rate + num_layers (:obj:`int`, *optional*, defaults to 1): + Number of Resnet layer block + resnet_groups (:obj:`int`, *optional*, defaults to `32`): + The number of groups to use for the Resnet and Attention block group norm + num_attention_heads (:obj:`int`, *optional*, defaults to `1`): + Number of attention heads for each attention block + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int + dropout: float = 0.0 + num_layers: int = 1 + resnet_groups: int = 32 + num_attention_heads: int = 1 + dtype: jnp.dtype = jnp.float32 + + def setup(self): + resnet_groups = self.resnet_groups if self.resnet_groups is not None else min(self.in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + FlaxResnetBlock2D( + in_channels=self.in_channels, + out_channels=self.in_channels, + dropout=self.dropout, + groups=resnet_groups, + dtype=self.dtype, + ) + ] + + attentions = [] + + for _ in range(self.num_layers): + attn_block = FlaxAttentionBlock( + channels=self.in_channels, + num_head_channels=self.num_attention_heads, + num_groups=resnet_groups, + dtype=self.dtype, + ) + attentions.append(attn_block) + + res_block = FlaxResnetBlock2D( + in_channels=self.in_channels, + out_channels=self.in_channels, + dropout=self.dropout, + groups=resnet_groups, + dtype=self.dtype, + ) + resnets.append(res_block) + + self.resnets = resnets + self.attentions = attentions + + def __call__(self, hidden_states, deterministic=True): + hidden_states = self.resnets[0](hidden_states, deterministic=deterministic) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + hidden_states = attn(hidden_states) + hidden_states = resnet(hidden_states, deterministic=deterministic) + + return hidden_states + + +class FlaxEncoder(nn.Module): + r""" + Flax Implementation of VAE Encoder. + + This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) + subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to + general usage and behavior. + + Finally, this model supports inherent JAX features such as: + - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) + - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) + - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) + - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) + + Parameters: + in_channels (:obj:`int`, *optional*, defaults to 3): + Input channels + out_channels (:obj:`int`, *optional*, defaults to 3): + Output channels + down_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(DownEncoderBlock2D)`): + DownEncoder block type + block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`): + Tuple containing the number of output channels for each block + layers_per_block (:obj:`int`, *optional*, defaults to `2`): + Number of Resnet layer for each block + norm_num_groups (:obj:`int`, *optional*, defaults to `32`): + norm num group + act_fn (:obj:`str`, *optional*, defaults to `silu`): + Activation function + double_z (:obj:`bool`, *optional*, defaults to `False`): + Whether to double the last output channels + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + Parameters `dtype` + """ + + in_channels: int = 3 + out_channels: int = 3 + down_block_types: Tuple[str] = ("DownEncoderBlock2D",) + block_out_channels: Tuple[int] = (64,) + layers_per_block: int = 2 + norm_num_groups: int = 32 + act_fn: str = "silu" + double_z: bool = False + dtype: jnp.dtype = jnp.float32 + + def setup(self): + block_out_channels = self.block_out_channels + # in + self.conv_in = nn.Conv( + block_out_channels[0], + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + # downsampling + down_blocks = [] + output_channel = block_out_channels[0] + for i, _ in enumerate(self.down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = FlaxDownEncoderBlock2D( + in_channels=input_channel, + out_channels=output_channel, + num_layers=self.layers_per_block, + resnet_groups=self.norm_num_groups, + add_downsample=not is_final_block, + dtype=self.dtype, + ) + down_blocks.append(down_block) + self.down_blocks = down_blocks + + # middle + self.mid_block = FlaxUNetMidBlock2D( + in_channels=block_out_channels[-1], + resnet_groups=self.norm_num_groups, + num_attention_heads=None, + dtype=self.dtype, + ) + + # end + conv_out_channels = 2 * self.out_channels if self.double_z else self.out_channels + self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) + self.conv_out = nn.Conv( + conv_out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + def __call__(self, sample, deterministic: bool = True): + # in + sample = self.conv_in(sample) + + # downsampling + for block in self.down_blocks: + sample = block(sample, deterministic=deterministic) + + # middle + sample = self.mid_block(sample, deterministic=deterministic) + + # end + sample = self.conv_norm_out(sample) + sample = nn.swish(sample) + sample = self.conv_out(sample) + + return sample + + +class FlaxDecoder(nn.Module): + r""" + Flax Implementation of VAE Decoder. + + This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) + subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to + general usage and behavior. + + Finally, this model supports inherent JAX features such as: + - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit) + - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) + - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap) + - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap) + + Parameters: + in_channels (:obj:`int`, *optional*, defaults to 3): + Input channels + out_channels (:obj:`int`, *optional*, defaults to 3): + Output channels + up_block_types (:obj:`Tuple[str]`, *optional*, defaults to `(UpDecoderBlock2D)`): + UpDecoder block type + block_out_channels (:obj:`Tuple[str]`, *optional*, defaults to `(64,)`): + Tuple containing the number of output channels for each block + layers_per_block (:obj:`int`, *optional*, defaults to `2`): + Number of Resnet layer for each block + norm_num_groups (:obj:`int`, *optional*, defaults to `32`): + norm num group + act_fn (:obj:`str`, *optional*, defaults to `silu`): + Activation function + double_z (:obj:`bool`, *optional*, defaults to `False`): + Whether to double the last output channels + dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): + parameters `dtype` + """ + + in_channels: int = 3 + out_channels: int = 3 + up_block_types: Tuple[str] = ("UpDecoderBlock2D",) + block_out_channels: int = (64,) + layers_per_block: int = 2 + norm_num_groups: int = 32 + act_fn: str = "silu" + dtype: jnp.dtype = jnp.float32 + + def setup(self): + block_out_channels = self.block_out_channels + + # z to block_in + self.conv_in = nn.Conv( + block_out_channels[-1], + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + # middle + self.mid_block = FlaxUNetMidBlock2D( + in_channels=block_out_channels[-1], + resnet_groups=self.norm_num_groups, + num_attention_heads=None, + dtype=self.dtype, + ) + + # upsampling + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + up_blocks = [] + for i, _ in enumerate(self.up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + + is_final_block = i == len(block_out_channels) - 1 + + up_block = FlaxUpDecoderBlock2D( + in_channels=prev_output_channel, + out_channels=output_channel, + num_layers=self.layers_per_block + 1, + resnet_groups=self.norm_num_groups, + add_upsample=not is_final_block, + dtype=self.dtype, + ) + up_blocks.append(up_block) + prev_output_channel = output_channel + + self.up_blocks = up_blocks + + # end + self.conv_norm_out = nn.GroupNorm(num_groups=self.norm_num_groups, epsilon=1e-6) + self.conv_out = nn.Conv( + self.out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + dtype=self.dtype, + ) + + def __call__(self, sample, deterministic: bool = True): + # z to block_in + sample = self.conv_in(sample) + + # middle + sample = self.mid_block(sample, deterministic=deterministic) + + # upsampling + for block in self.up_blocks: + sample = block(sample, deterministic=deterministic) + + sample = self.conv_norm_out(sample) + sample = nn.swish(sample) + sample = self.conv_out(sample) + + return sample + + + +def load_pretrained_vae(modelname: str, revision: str = "bf16") -> dict: + """Load a pretrained Flax AutoencoderKL config and params from the HF Hub. + + Works with both full Stable Diffusion pipeline repos (VAE under the `vae` + subfolder, flax weights usually on the `flax`/`bf16` revisions) and + standalone VAE repos (files at the repo root). + Returns a dict with `config` (dict) and `params` (nested param tree). + """ + from flax.serialization import msgpack_restore + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError, RevisionNotFoundError + + candidates = [ + {"revision": revision, "subfolder": "vae"}, + {"revision": "flax", "subfolder": "vae"}, + {"revision": revision}, + {}, + ] + last_error = None + for candidate in candidates: + try: + config_path = hf_hub_download(modelname, "config.json", **candidate) + weights_path = hf_hub_download(modelname, "diffusion_flax_model.msgpack", **candidate) + break + except (EntryNotFoundError, RevisionNotFoundError) as e: + last_error = e + else: + raise FileNotFoundError( + f"No flax VAE weights (diffusion_flax_model.msgpack) found in {modelname}" + ) from last_error + + with open(config_path) as f: + config = json.load(f) + with open(weights_path, "rb") as f: + params = msgpack_restore(f.read()) + return {"config": config, "params": params} diff --git a/pyproject.toml b/pyproject.toml index 9c178e2..6d1b8a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,8 +25,23 @@ dependencies = [ "albumentations", "rich", "python-dotenv", + "transformers", + "huggingface_hub", + "wandb", + # diffusers removed all Flax support upstream (huggingface/diffusers#14169); + # only needed for the diffusers_unet_simple architecture, the VAE is vendored + "diffusers>=0.29,<0.35", ] license = "MIT" +[project.optional-dependencies] +test = ["pytest"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +markers = [ + "network: tests that download pretrained weights from the HuggingFace Hub", +] + [tool.setuptools.packages.find] include = ["flaxdiff*"] \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..d9c41fd --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,19 @@ +import os + +# Tests must run identically on any machine, CPU is enough +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +import jax +import jax.numpy as jnp +import pytest + + +@pytest.fixture +def rng(): + return jax.random.PRNGKey(0) + + +@pytest.fixture +def text_context(): + # Shape of the default CLIP-L/14 text context, no need for the actual encoder + return jnp.ones((2, 77, 768), dtype=jnp.float32) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..b7ef699 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,77 @@ +"""Train -> wandb config -> inference reconstruction round-trip tests. + +parse_config must rebuild exactly the model that was trained. Any key that +is silently dropped or remapped means the inference pipeline runs a +different model than the one in the checkpoint. +""" + +import jax +import jax.numpy as jnp +import pytest + +from flaxdiff.inference.utils import parse_config + + +def make_config(model_overrides=None, arguments_overrides=None): + """A minimal wandb-style config, the same shape training.py logs.""" + model = { + "emb_features": 64, + "dtype": "bfloat16", + "precision": "high", + "activation": "swish", + "output_channels": 3, + "norm_groups": 4, + "patch_size": 4, + "num_layers": 2, + "num_heads": 2, + "dropout_rate": 0.0, + "mlp_ratio": 2, + "use_hilbert": False, + "use_zigzag": False, + } + model.update(model_overrides or {}) + arguments = { + "architecture": "simple_dit", + "image_size": 32, + "noise_schedule": "edm", + } + arguments.update(arguments_overrides or {}) + return { + "model": model, + "architecture": arguments["architecture"], + "arguments": arguments, + "input_config": { + "sample_data_key": "image", + "sample_data_shape": (32, 32, 3), + "conditions": [], + }, + } + + +def test_parse_config_rebuilds_model(): + result = parse_config(make_config()) + model = result["model"] + assert type(model).__name__ == "SimpleDiT" + assert model.emb_features == 64 + assert model.num_layers == 2 + assert model.dtype == jnp.bfloat16 + assert model.precision == jax.lax.Precision.HIGH + assert model.activation is jax.nn.swish + + +def test_parse_config_noise_schedule_selection(): + edm = parse_config(make_config()) + assert type(edm["noise_schedule"]).__name__ == "KarrasVENoiseScheduler" + + cosine = parse_config(make_config(arguments_overrides={"noise_schedule": "cosine"})) + assert type(cosine["noise_schedule"]).__name__ == "CosineNoiseScheduler" + + +@pytest.mark.xfail(strict=True, reason="bug: parse_config silently drops any string value containing a dot") +def test_parse_config_preserves_dotted_values(): + """String config values with a dot in them (module paths, filenames, version + strings) must survive the round trip, not vanish into class defaults.""" + result = parse_config(make_config(model_overrides={"activation": "jax.nn.mish"})) + model = result["model"] + # a dropped key falls back to the class default (swish) instead of erroring + assert model.activation is jax.nn.mish, "activation was silently dropped" diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..c86f173 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,117 @@ +"""Forward-pass tests for every model architecture. + +Every architecture must build, run a forward pass at a small resolution, +and return the input shape back in float32. This is the first line of +defense against dead configs and shape bugs. +""" + +import jax +import jax.numpy as jnp +import pytest + +from flaxdiff.models.simple_unet import Unet +from flaxdiff.models.simple_dit import SimpleDiT +from flaxdiff.models.simple_mmdit import SimpleMMDiT +from flaxdiff.models.simple_vit import UViT +from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT + +RES = 32 + + +def run_forward(model, rng, x, temb, textcontext): + params = model.init(rng, x, temb, textcontext) + return model.apply(params, x, temb, textcontext) + + +def small_inputs(rng, res=RES, channels=3): + x = jax.random.normal(rng, (2, res, res, channels)) + temb = jnp.ones((2,)) + textcontext = jnp.ones((2, 77, 768), dtype=jnp.float32) + return x, temb, textcontext + + +def test_unet_forward(rng): + model = Unet( + emb_features=64, + feature_depths=[16, 32], + attention_configs=[None, {"heads": 2, "dtype": jnp.float32, "flash_attention": False, + "use_projection": False, "use_self_and_cross": False}], + num_res_blocks=1, + num_middle_res_blocks=1, + ) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + assert out.dtype == jnp.float32 + + +def test_simple_dit_forward(rng): + model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + assert out.dtype == jnp.float32 + + +@pytest.mark.xfail(strict=True, reason="bug: DiT final_proj keeps compute dtype, bf16 output vs fp32 loss") +def test_simple_dit_bf16_outputs_fp32(rng): + model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, + mlp_ratio=2, dtype=jnp.bfloat16) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.dtype == jnp.float32 + + +@pytest.mark.xfail(strict=True, reason="bug: unpatchify assumes a square patch grid") +def test_simple_dit_non_square(rng): + model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x = jax.random.normal(rng, (2, 16, 64, 3)) + temb = jnp.ones((2,)) + textcontext = jnp.ones((2, 77, 768), dtype=jnp.float32) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + + +def test_simple_mmdit_forward(rng): + model = SimpleMMDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2, mlp_ratio=2) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + + +def test_uvit_forward(rng): + model = UViT(patch_size=4, emb_features=64, num_layers=4, num_heads=2) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + + +@pytest.mark.parametrize("kwargs", [ + {}, + {"use_hilbert": True}, + {"use_zigzag": True}, + {"use_zigzag": True, "use_2d_fusion": True}, + {"use_hilbert": True, "use_2d_fusion": True}, +]) +def test_hybrid_dit_forward(rng, kwargs): + model = HybridSSMAttentionDiT(patch_size=4, emb_features=64, num_layers=4, + num_heads=2, mlp_ratio=2, ssm_state_dim=8, **kwargs) + x, temb, textcontext = small_inputs(rng) + out = run_forward(model, rng, x, temb, textcontext) + assert out.shape == x.shape + + +@pytest.mark.xfail(strict=True, reason="bug: UViT applies the hilbert permutation twice") +def test_uvit_hilbert_matches_raster_information(rng): + """A zero-layer sanity check: with the permutation applied and inverted once, + patch content must land back at its own spatial position. UViT permutes twice + on the way in but inverts once on the way out, scrambling the output.""" + from flaxdiff.models.hilbert import hilbert_patchify, hilbert_unpatchify, hilbert_indices, inverse_permutation + + x = jax.random.normal(rng, (1, 16, 16, 3)) + patches, inv_idx = hilbert_patchify(x, 4) + # UViT's forward applies idx again on the already-permuted patches + idx = hilbert_indices(4, 4) + double_permuted = patches[:, idx, :] + rec = hilbert_unpatchify(double_permuted, inv_idx, 4, 16, 16, 3) + assert jnp.allclose(rec, x) diff --git a/tests/test_predictors.py b/tests/test_predictors.py new file mode 100644 index 0000000..31a7ec8 --- /dev/null +++ b/tests/test_predictors.py @@ -0,0 +1,67 @@ +"""Round-trip tests for the prediction transforms (parameterizations). + +forward_diffusion produces (x_t, c_in, target); backward_diffusion must +recover x0 and epsilon exactly from that target. This must hold for every +parameterization on both VP and VE schedules. +""" + +import jax +import jax.numpy as jnp +import pytest + +from flaxdiff.predictors import ( + EpsilonPredictionTransform, + DirectPredictionTransform, + VPredictionTransform, + KarrasPredictionTransform, +) +from flaxdiff.schedulers import CosineNoiseScheduler, KarrasVENoiseScheduler +from flaxdiff.schedulers.common import get_coeff_shapes_tuple + +TRANSFORMS = [ + ("epsilon", EpsilonPredictionTransform()), + ("x0", DirectPredictionTransform()), + ("v", VPredictionTransform()), + ("karras", KarrasPredictionTransform(sigma_data=0.5)), +] +SCHEDULES = [ + ("cosine", CosineNoiseScheduler(1000), jnp.array([10, 300, 600, 900])), + ("karras_ve", KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5), jnp.array([0.2, 0.4, 0.6, 0.8])), +] + + +@pytest.mark.parametrize("tname,transform", TRANSFORMS) +@pytest.mark.parametrize("sname,schedule,steps", SCHEDULES) +def test_forward_backward_roundtrip(tname, transform, sname, schedule, steps, rng): + key0, key1 = jax.random.split(rng) + # one timestep per batch element + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + rates = schedule.get_rates(steps, get_coeff_shapes_tuple(x0)) + + xt, c_in, target = transform.forward_diffusion(x0, noise, rates) + # A model that outputs the exact target must recover x0 and noise + recovered_x0, recovered_noise = transform.backward_diffusion(xt, target, rates) + + assert jnp.max(jnp.abs(recovered_x0 - x0)) < 5e-3, f"{tname} on {sname}: x0 not recovered" + assert jnp.max(jnp.abs(recovered_noise - noise)) < 5e-3, f"{tname} on {sname}: noise not recovered" + + +def test_karras_preconditioning_matches_paper(rng): + """c_in, c_skip and c_out must match Karras et al. 2022 Table 1.""" + transform = KarrasPredictionTransform(sigma_data=0.5) + schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + steps = jnp.array([0.3]) + rates = schedule.get_rates(steps, (-1, 1, 1, 1)) + _, sigma = rates + sd = 0.5 + + c_in = transform.get_input_scale(rates) + assert jnp.allclose(c_in, 1 / jnp.sqrt(sigma**2 + sd**2), rtol=1e-4) + + x_t = jax.random.normal(rng, (1, 8, 8, 3)) + raw = jax.random.normal(jax.random.fold_in(rng, 1), (1, 8, 8, 3)) + c_skip = sd**2 / (sd**2 + sigma**2) + c_out = sigma * sd / jnp.sqrt(sd**2 + sigma**2) + expected = c_skip * x_t + c_out * raw + assert jnp.allclose(transform.pred_transform(x_t, raw, rates), expected, rtol=1e-4) diff --git a/tests/test_samplers.py b/tests/test_samplers.py new file mode 100644 index 0000000..0e9eac0 --- /dev/null +++ b/tests/test_samplers.py @@ -0,0 +1,208 @@ +"""End-to-end sampler tests against an analytic oracle denoiser. + +For gaussian data x0 ~ N(0, s^2 I), the optimal denoiser has a closed form: +E[x0 | x_t] = alpha * s^2 / (alpha^2 s^2 + sigma^2) * x_t. A sampler +integrating the reverse process with this oracle must produce samples with +mean 0 and std s. This catches integrator bugs, step-domain bugs, and +shape bugs in one place, with no training involved. +""" + +import jax +import jax.numpy as jnp +import pytest +from flax import linen as nn + +from flaxdiff.inputs import DiffusionInputConfig +from flaxdiff.predictors import EpsilonPredictionTransform, KarrasPredictionTransform +from flaxdiff.samplers.ddim import DDIMSampler +from flaxdiff.samplers.ddpm import DDPMSampler, SimpleDDPMSampler +from flaxdiff.samplers.euler import EulerSampler, EulerAncestralSampler +from flaxdiff.samplers.heun_sampler import HeunSampler +from flaxdiff.samplers.multistep_dpm import MultiStepDPM +from flaxdiff.samplers.rk4_sampler import RK4Sampler +from flaxdiff.schedulers import CosineNoiseScheduler, KarrasVENoiseScheduler +from flaxdiff.schedulers.common import get_coeff_shapes_tuple +from flaxdiff.utils import RandomMarkovState + +DATA_STD = 0.3 + +input_config = DiffusionInputConfig( + sample_data_key="image", + sample_data_shape=(8, 8, 3), + conditions=[], +) + + +class VPOracle(nn.Module): + """Optimal epsilon-predictor for x0 ~ N(0, DATA_STD^2) on a VP schedule. + + The sampler feeds raw timesteps for discrete schedules, so rates are + recomputed from the schedule inside the module. + """ + schedule: CosineNoiseScheduler + + @nn.compact + def __call__(self, x, temb): + alpha, sigma = self.schedule.get_rates(temb, get_coeff_shapes_tuple(x)) + return x * sigma / (alpha**2 * DATA_STD**2 + sigma**2) + + +class KarrasOracle(nn.Module): + """Optimal raw-F predictor under the Karras preconditioning (VE, alpha=1). + + Receives x * c_in and temb = log(sigma)/4; must output F such that + c_skip * x + c_out * F = E[x0 | x] = s^2/(s^2 + sigma^2) * x. + """ + sigma_data: float = 0.5 + + @nn.compact + def __call__(self, x_scaled, temb): + sigma = jnp.exp(4.0 * temb).reshape(get_coeff_shapes_tuple(x_scaled)) + sd = self.sigma_data + c_in = 1 / jnp.sqrt(sigma**2 + sd**2) + c_skip = sd**2 / (sd**2 + sigma**2) + c_out = sigma * sd / jnp.sqrt(sd**2 + sigma**2) + x = x_scaled / c_in + x0 = DATA_STD**2 / (DATA_STD**2 + sigma**2) * x + return (x0 - c_skip * x) / c_out + + +def make_vp_sampler(sampler_class, **kwargs): + schedule = CosineNoiseScheduler(1000) + model = VPOracle(schedule=schedule) + return model, sampler_class( + model=model, + noise_schedule=schedule, + model_output_transform=EpsilonPredictionTransform(), + input_config=input_config, + guidance_scale=0.0, + **kwargs, + ) + + +def make_karras_sampler(sampler_class, **kwargs): + schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + model = KarrasOracle() + return model, sampler_class( + model=model, + noise_schedule=schedule, + model_output_transform=KarrasPredictionTransform(sigma_data=0.5), + input_config=input_config, + guidance_scale=0.0, + **kwargs, + ) + + +def assert_gaussian_stats(samples, std=DATA_STD, tol=0.05): + assert jnp.all(jnp.isfinite(samples)), "sampler produced non-finite values" + assert abs(float(jnp.mean(samples))) < tol + assert abs(float(jnp.std(samples)) - std) < tol + + +def generate(model, sampler, **kwargs): + params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,))) + defaults = dict( + num_samples=256, + resolution=8, + diffusion_steps=100, + start_step=1000, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + ) + defaults.update(kwargs) + return sampler.generate_samples(params, **defaults) + + +@pytest.mark.parametrize("sampler_class", [EulerSampler, DDIMSampler, SimpleDDPMSampler]) +def test_vp_sampler_converges(sampler_class): + model, sampler = make_vp_sampler(sampler_class) + samples = generate(model, sampler) + assert_gaussian_stats(samples) + + +@pytest.mark.parametrize("sampler_class", [EulerSampler, EulerAncestralSampler, DDIMSampler, HeunSampler, MultiStepDPM]) +def test_karras_sampler_converges(sampler_class): + model, sampler = make_karras_sampler(sampler_class) + samples = generate(model, sampler) + assert_gaussian_stats(samples) + + +@pytest.mark.xfail(strict=True, reason="bug: default start_step uses max_timesteps but the step domain is hardcoded to [0, 1000]") +def test_karras_sampler_default_start_step(): + """With no explicit start_step the sampler must still denoise from sigma_max. + Today max_timesteps=1 collapses the schedule to a single no-op step and + the 'samples' are just the initial noise.""" + model, sampler = make_karras_sampler(EulerSampler) + samples = generate(model, sampler, start_step=None) + assert_gaussian_stats(samples) + + +@pytest.mark.xfail(strict=True, reason="bug: sample_model reshapes rates with the default (-1,1,1,1), breaking 5D video shapes for every sampler") +def test_ddim_video_samples(): + """DDIM's own take_next_step is 5D-safe, but the shared sample_model wrapper + still scales the input with 4D-reshaped rates, so video is broken everywhere.""" + model, sampler = make_karras_sampler(DDIMSampler) + params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,))) + samples = sampler.generate_samples( + params, + num_samples=2, + resolution=8, + sequence_length=3, + diffusion_steps=50, + start_step=1000, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + ) + assert samples.shape == (2, 3, 8, 8, 3) + assert jnp.all(jnp.isfinite(samples)) + + +@pytest.mark.xfail(strict=True, reason="bug: samplers reshape rates with the default (-1,1,1,1), breaking 5D video shapes") +def test_euler_video_samples(): + model, sampler = make_karras_sampler(EulerSampler) + params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,))) + samples = sampler.generate_samples( + params, + num_samples=2, + resolution=8, + sequence_length=3, + diffusion_steps=50, + start_step=1000, + rngstate=RandomMarkovState(jax.random.PRNGKey(2)), + ) + assert samples.shape == (2, 3, 8, 8, 3) + + +@pytest.mark.xfail(strict=True, reason="bug: DDPMSampler casts the step array through int(), crashing for batch > 1") +def test_ddpm_sampler_converges(): + model, sampler = make_vp_sampler(DDPMSampler) + samples = generate(model, sampler, diffusion_steps=1000) + assert_gaussian_stats(samples) + + +@pytest.mark.xfail(strict=True, reason="bug: DDIM eta > 0 calls .sqrt() on a jax array and uses the wrong variance split") +def test_ddim_eta_converges(): + model, sampler = make_vp_sampler(DDIMSampler, eta=0.5) + samples = generate(model, sampler) + assert_gaussian_stats(samples) + + +@pytest.mark.xfail(strict=True, reason="bug: RK4Sampler jits over its python-callable argument") +def test_rk4_sampler_runs(): + model, sampler = make_karras_sampler(RK4Sampler) + samples = generate(model, sampler, diffusion_steps=20) + assert jnp.all(jnp.isfinite(samples)) + + +@pytest.mark.xfail(strict=True, reason="bug: MultiStepDPM keeps stale derivative history across generate calls") +def test_multistep_dpm_reentrant(): + model, sampler = make_karras_sampler(MultiStepDPM) + first = generate(model, sampler) + second = generate(model, sampler) + assert jnp.allclose(first, second, atol=1e-5) + + +@pytest.mark.parametrize("spacing", ["quadratic", "karras", "exponential"]) +@pytest.mark.xfail(strict=True, reason="bug: non-linear timestep spacings produce log(0) or duplicate int16 steps (dt=0)") +def test_timestep_spacing_produces_finite_samples(spacing): + model, sampler = make_karras_sampler(EulerSampler, timestep_spacing=spacing) + samples = generate(model, sampler, diffusion_steps=50) + assert_gaussian_stats(samples, tol=0.1) diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py new file mode 100644 index 0000000..b54cc54 --- /dev/null +++ b/tests/test_schedulers.py @@ -0,0 +1,87 @@ +"""Invariant tests for the noise schedulers. + +These encode the properties the rest of the library relies on: +variance preservation for VP schedules, alpha=1 for the generalized (VE) +schedules, and exact invertibility of the forward diffusion. +""" + +import jax +import jax.numpy as jnp +import pytest + +from flaxdiff.schedulers import ( + CosineNoiseScheduler, + LinearNoiseSchedule, + KarrasVENoiseScheduler, + EDMNoiseScheduler, +) +from flaxdiff.schedulers.common import get_coeff_shapes_tuple + +VP_SCHEDULES = [ + ("cosine", lambda: CosineNoiseScheduler(1000), jnp.array([10, 300, 600, 900])), + ("linear", lambda: LinearNoiseSchedule(1000), jnp.array([10, 300, 600, 900])), +] +VE_SCHEDULES = [ + ("karras_ve", lambda: KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5), jnp.array([0.05, 0.3, 0.6, 0.95])), + ("edm", lambda: EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5), jnp.array([0.05, 0.3, 0.6, 0.95])), +] + + +@pytest.mark.parametrize("name,make,steps", VP_SCHEDULES) +def test_vp_variance_preserving(name, make, steps): + schedule = make() + alpha, sigma = schedule.get_rates(steps, shape=(-1,)) + assert jnp.allclose(alpha**2 + sigma**2, 1.0, atol=1e-5) + + +@pytest.mark.parametrize("name,make,steps", VE_SCHEDULES) +def test_ve_signal_rate_is_one(name, make, steps): + schedule = make() + alpha, sigma = schedule.get_rates(steps, shape=(-1,)) + assert jnp.allclose(alpha, 1.0) + # Noise grows monotonically with t + assert jnp.all(jnp.diff(sigma) > 0) + + +@pytest.mark.parametrize("name,make,steps", VP_SCHEDULES + VE_SCHEDULES) +def test_forward_diffusion_invertible(name, make, steps, rng): + schedule = make() + key0, key1 = jax.random.split(rng) + x0 = jax.random.normal(key0, (4, 8, 8, 3)) + noise = jax.random.normal(key1, (4, 8, 8, 3)) + xt = schedule.add_noise(x0, noise, steps) + recovered = schedule.remove_all_noise(xt, noise, steps) + assert jnp.max(jnp.abs(recovered - x0)) < 1e-4 + + +def test_karras_weights_match_edm_lambda(): + """KarrasVE loss weights must equal EDM's lambda(sigma) = (s^2+sd^2)/(s*sd)^2. + The current epsilon guard halves the weight at sigma_min.""" + schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + # steps away from sigma_min; the sigma_min case is the xfail test below + steps = jnp.array([0.3, 0.6, 1.0]) + sigma = schedule.get_sigmas(steps) + expected = (sigma**2 + 0.5**2) / ((sigma * 0.5) ** 2) + got = schedule.get_weights(steps, shape=(-1,)) + assert jnp.allclose(got, expected, rtol=1e-3) + + +@pytest.mark.xfail(strict=True, reason="bug: epsilon guard distorts the weight near sigma_min") +def test_karras_weights_at_sigma_min(): + schedule = KarrasVENoiseScheduler(1, sigma_min=0.002, sigma_max=80, rho=7, sigma_data=0.5) + sigma = jnp.array([0.002]) + expected = (sigma**2 + 0.5**2) / ((sigma * 0.5) ** 2) + # steps=0 maps to sigma_min under the karras rho spacing + got = schedule.get_weights(jnp.array([0.0]), shape=(-1,)) + assert jnp.allclose(got, expected, rtol=1e-2) + + +def test_edm_lognormal_sigma_distribution(rng): + """EDM training sigmas must follow exp(N(-1.2, 1.2^2)) when timesteps=1.""" + from flaxdiff.utils import RandomMarkovState + + schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5) + steps, _ = schedule.generate_timesteps(20000, RandomMarkovState(rng)) + log_sigma = jnp.log(schedule.get_sigmas(steps)) + assert abs(float(jnp.mean(log_sigma)) - (-1.2)) < 0.05 + assert abs(float(jnp.std(log_sigma)) - 1.2) < 0.05 diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 0000000..f8d86f6 --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,70 @@ +"""Trainer smoke tests: training without wandb, and checkpoint save/restore. + +These run on CPU with a tiny regression model through SimpleTrainer, which +owns the state/checkpoint/loop machinery shared by all trainers. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import pytest +from flax import linen as nn + +from flaxdiff.trainer.simple_trainer import SimpleTrainer + + +class TinyModel(nn.Module): + @nn.compact + def __call__(self, x): + return nn.Dense(2)(x) + + +def make_trainer(tmp_path, name="smoke", load_from_checkpoint=None): + return SimpleTrainer( + model=TinyModel(), + input_shapes={"x": (4,)}, + optimizer=optax.adam(1e-3), + rngs=jax.random.PRNGKey(0), + name=name, + wandb_config=None, + distributed_training=False, + checkpoint_base_path=str(tmp_path), + load_from_checkpoint=load_from_checkpoint, + ) + + +def batch_iterator(): + while True: + yield {"image": jnp.ones((8, 4)), "label": jnp.zeros((8, 2))} + + +@pytest.mark.xfail(strict=True, reason="bug: self.wandb is only assigned when wandb_config is set, training without wandb crashes") +def test_fit_without_wandb(tmp_path): + trainer = make_trainer(tmp_path) + data = {"train": batch_iterator, "train_len": 32} + trainer.fit(data, train_steps_per_epoch=4, epochs=1, val_steps_per_epoch=0) + + +def test_save_writes_checkpoint(tmp_path): + """save() swallows exceptions, so assert the checkpoint actually landed.""" + trainer = make_trainer(tmp_path) + trainer.save(epoch=0, step=1) + assert trainer.checkpointer.latest_step() == 1 + + +@pytest.mark.xfail(strict=True, reason="bug: restore rebuilds the TrainState with tx.init, discarding opt_state and the step counter") +def test_restore_preserves_optimizer_state(tmp_path): + trainer = make_trainer(tmp_path) + + # One real update so the adam moments and step counter are non-trivial + grads = jax.tree.map(jnp.ones_like, trainer.state.params) + trainer.state = trainer.state.apply_gradients(grads=grads) + trainer.save(epoch=0, step=1) + + restored = make_trainer(tmp_path, load_from_checkpoint=trainer.checkpoint_path()) + assert int(restored.state.step) == 1, "optimizer step counter was reset" + + old_mu = jax.tree.leaves(trainer.state.opt_state) + new_mu = jax.tree.leaves(restored.state.opt_state) + assert all(np.allclose(a, b) for a, b in zip(old_mu, new_mu)), "adam moments were reset" diff --git a/tests/test_vae.py b/tests/test_vae.py new file mode 100644 index 0000000..96904ec --- /dev/null +++ b/tests/test_vae.py @@ -0,0 +1,33 @@ +"""Vendored Stable Diffusion VAE tests. + +Marked as network tests: they download the pretrained weights from the +HuggingFace Hub on first run. Excluded in CI (-m "not network"). +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +pytestmark = pytest.mark.network + + +@pytest.fixture(scope="module") +def vae(): + from flaxdiff.models.autoencoder.diffusers import StableDiffusionVAE + return StableDiffusionVAE(dtype=jnp.float32) + + +def test_vae_shapes(vae): + assert vae.downscale_factor == 8 + assert vae.latent_channels == 4 + + +def test_vae_roundtrip_reconstructs(vae, rng): + # A smooth image should survive the encode/decode roundtrip well + ramp = jnp.linspace(-0.8, 0.8, 64) + x = jnp.broadcast_to(ramp[None, :, None, None], (1, 64, 64, 3)).transpose(0, 2, 1, 3) + rec = vae.decode(vae.encode(x)) + mse = float(jnp.mean((rec - x) ** 2)) + psnr = 10 * np.log10(4.0 / mse) + assert psnr > 20, f"reconstruction too poor: {psnr:.1f}dB"