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

flaxdiff/metrics/utils.py

@@ -0,0 +1,49 @@
+# Mostly derived from
+# https://github.com/matthias-wright/jax-fid
+
+import jax
+import flax
+import numpy as np
+from tqdm import tqdm
+import requests
+import os
+import tempfile
+
+
+def download(url, ckpt_dir=None):
+    name = url[url.rfind('/') + 1 : url.rfind('?')]
+    if ckpt_dir is None:
+        ckpt_dir = tempfile.gettempdir()
+    ckpt_dir = os.path.join(ckpt_dir, 'jax_fid')
+    ckpt_file = os.path.join(ckpt_dir, name)
+    if not os.path.exists(ckpt_file):
+        print(f'Downloading: \"{url[:url.rfind("?")]}\" to {ckpt_file}')
+        if not os.path.exists(ckpt_dir): 
+            os.makedirs(ckpt_dir)
+
+        response = requests.get(url, stream=True)
+        total_size_in_bytes = int(response.headers.get('content-length', 0))
+        progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+        
+        # first create temp file, in case the download fails
+        ckpt_file_temp = os.path.join(ckpt_dir, name + '.temp')
+        with open(ckpt_file_temp, 'wb') as file:
+            for data in response.iter_content(chunk_size=1024):
+                progress_bar.update(len(data))
+                file.write(data)
+        progress_bar.close()
+        
+        if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+            print('An error occured while downloading, please try again.')
+            if os.path.exists(ckpt_file_temp):
+                os.remove(ckpt_file_temp)
+        else:
+            # if download was successful, rename the temp file
+            os.rename(ckpt_file_temp, ckpt_file)
+    return ckpt_file
+
+
+def get(dictionary, key):
+    if dictionary is None or key not in dictionary:
+        return None
+    return dictionary[key]

flaxdiff/models/__init__.py

@@ -0,0 +1 @@
+from .simple_unet import *

flaxdiff/models/attention.py

@@ -0,0 +1,368 @@
+"""
+Some Code ported from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_flax.py
+"""
+
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Dict, Callable, Sequence, Any, Union, Tuple, Optional
+from flax.typing import Dtype, PrecisionLike
+import einops
+import functools
+import math
+from .common import kernel_init
+import jax.experimental.pallas.ops.tpu.flash_attention
+
+class EfficientAttention(nn.Module):
+    """
+    Based on the pallas attention implementation.
+    """
+    query_dim: int
+    heads: int = 4
+    dim_head: int = 64
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_bias: bool = True
+    kernel_init: Callable = kernel_init(1.0)
+    force_fp32_for_softmax: bool = True
+
+    def setup(self):
+        inner_dim = self.dim_head * self.heads
+        # Weights were exported with old names {to_q, to_k, to_v, to_out}
+        dense = functools.partial(
+            nn.Dense,
+            self.heads * self.dim_head,
+            precision=self.precision, 
+            use_bias=self.use_bias, 
+            kernel_init=self.kernel_init, 
+            dtype=self.dtype
+        )
+        self.query = dense(name="to_q")
+        self.key = dense(name="to_k")
+        self.value = dense(name="to_v")
+        
+        self.proj_attn = nn.DenseGeneral(self.query_dim, use_bias=False, precision=self.precision, 
+                                     kernel_init=self.kernel_init, dtype=self.dtype, name="to_out_0")
+        # self.attnfn = make_fast_generalized_attention(qkv_dim=inner_dim, lax_scan_unroll=16)
+    
+    def _reshape_tensor_to_head_dim(self, tensor):
+        batch_size, _, seq_len, dim = tensor.shape
+        head_size = self.heads
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
+        return tensor
+    
+    def _reshape_tensor_from_head_dim(self, tensor):
+        batch_size, _, seq_len, dim = tensor.shape
+        head_size = self.heads
+        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
+        tensor = tensor.reshape(batch_size, 1, seq_len, dim * head_size)
+        return tensor
+
+    @nn.compact
+    def __call__(self, x:jax.Array, context=None):
+        # print(x.shape)
+        # x has shape [B, H * W, C]
+        context = x if context is None else context
+        
+        orig_x_shape = x.shape
+        if len(x.shape) == 4:
+            B, H, W, C = x.shape
+            x = x.reshape((B, 1, H * W, C))
+        else:
+            B, SEQ, C = x.shape
+            x = x.reshape((B, 1, SEQ, C))
+        
+        if len(context.shape) == 4:
+            B, _H, _W, _C = context.shape
+            context = context.reshape((B, 1, _H * _W, _C))
+        else:
+            B, SEQ, _C = context.shape
+            context = context.reshape((B, 1, SEQ, _C))
+        
+        query = self.query(x)
+        key = self.key(context)
+        value = self.value(context)
+        
+        query = self._reshape_tensor_to_head_dim(query)
+        key = self._reshape_tensor_to_head_dim(key)
+        value = self._reshape_tensor_to_head_dim(value)
+        
+        hidden_states = jax.experimental.pallas.ops.tpu.flash_attention.flash_attention(
+            query, key, value, None
+        )
+        
+        hidden_states = self._reshape_tensor_from_head_dim(hidden_states)
+        
+        
+        # hidden_states = nn.dot_product_attention(
+        #     query, key, value, dtype=self.dtype, broadcast_dropout=False, dropout_rng=None, precision=self.precision
+        # )
+        
+        proj = self.proj_attn(hidden_states)
+        
+        proj = proj.reshape(orig_x_shape)
+        
+        return proj
+
+class NormalAttention(nn.Module):
+    """
+    Simple implementation of the normal attention.
+    """
+    query_dim: int
+    heads: int = 4
+    dim_head: int = 64
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_bias: bool = True
+    kernel_init: Callable = kernel_init(1.0)
+    force_fp32_for_softmax: bool = True
+
+    def setup(self):
+        inner_dim = self.dim_head * self.heads
+        dense = functools.partial(
+            nn.DenseGeneral,
+            features=[self.heads, self.dim_head], 
+            axis=-1, 
+            precision=self.precision, 
+            use_bias=self.use_bias, 
+            kernel_init=self.kernel_init, 
+            dtype=self.dtype
+        )
+        self.query = dense(name="to_q")
+        self.key = dense(name="to_k")
+        self.value = dense(name="to_v")
+
+        self.proj_attn = nn.DenseGeneral(
+            self.query_dim, 
+            axis=(-2, -1), 
+            precision=self.precision, 
+            use_bias=self.use_bias, 
+            dtype=self.dtype, 
+            name="to_out_0",
+            kernel_init=self.kernel_init
+            # kernel_init=jax.nn.initializers.xavier_uniform()
+        )
+
+    @nn.compact
+    def __call__(self, x, context=None):
+        # x has shape [B, H, W, C]
+        orig_x_shape = x.shape
+        if len(x.shape) == 4:
+            B, H, W, C = x.shape
+            x = x.reshape((B, H*W, C))
+        context = x if context is None else context
+        if len(context.shape) == 4:
+            context = context.reshape((B, H*W, C))
+        query = self.query(x)
+        key = self.key(context)
+        value = self.value(context)
+        
+        hidden_states = nn.dot_product_attention(
+            query, key, value, dtype=self.dtype, broadcast_dropout=False, 
+            dropout_rng=None, precision=self.precision, force_fp32_for_softmax=self.force_fp32_for_softmax,
+            deterministic=True
+        )
+        proj = self.proj_attn(hidden_states)
+        proj = proj.reshape(orig_x_shape)
+        return proj
+
+class FlaxGEGLU(nn.Module):
+    r"""
+    Flax implementation of a Linear layer followed by the variant of the gated linear unit activation function from
+    https://arxiv.org/abs/2002.05202.
+
+    Parameters:
+        dim (:obj:`int`):
+            Input hidden states dimension
+        dropout (:obj:`float`, *optional*, defaults to 0.0):
+            Dropout rate
+        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
+            Parameters `dtype`
+    """
+
+    dim: int
+    dropout: float = 0.0
+    dtype: jnp.dtype = jnp.float32
+    precision: Any = jax.lax.Precision.DEFAULT
+
+    def setup(self):
+        inner_dim = self.dim * 4
+        self.proj = nn.Dense(inner_dim * 2, dtype=self.dtype, precision=self.precision)
+
+    def __call__(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_linear, hidden_gelu = jnp.split(hidden_states, 2, axis=-1)
+        return hidden_linear * nn.gelu(hidden_gelu)
+    
+class FlaxFeedForward(nn.Module):
+    r"""
+    Flax module that encapsulates two Linear layers separated by a non-linearity. It is the counterpart of PyTorch's
+    [`FeedForward`] class, with the following simplifications:
+    - The activation function is currently hardcoded to a gated linear unit from:
+    https://arxiv.org/abs/2002.05202
+    - `dim_out` is equal to `dim`.
+    - The number of hidden dimensions is hardcoded to `dim * 4` in [`FlaxGELU`].
+
+    Parameters:
+        dim (:obj:`int`):
+            Inner hidden states dimension
+        dropout (:obj:`float`, *optional*, defaults to 0.0):
+            Dropout rate
+        dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
+            Parameters `dtype`
+    """
+
+    dim: int
+    dtype: jnp.dtype = jnp.float32
+    precision: Any = jax.lax.Precision.DEFAULT
+
+    def setup(self):
+        # The second linear layer needs to be called
+        # net_2 for now to match the index of the Sequential layer
+        self.net_0 = FlaxGEGLU(self.dim, self.dtype, precision=self.precision)
+        self.net_2 = nn.Dense(self.dim, dtype=self.dtype, precision=self.precision)
+
+    def __call__(self, hidden_states):
+        hidden_states = self.net_0(hidden_states)
+        hidden_states = self.net_2(hidden_states)
+        return hidden_states
+
+class BasicTransformerBlock(nn.Module):
+    # Has self and cross attention
+    query_dim: int
+    heads: int = 4
+    dim_head: int = 64
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_bias: bool = True
+    kernel_init: Callable = kernel_init(1.0)
+    use_flash_attention:bool = False
+    use_cross_only:bool = False
+    only_pure_attention:bool = False
+    force_fp32_for_softmax: bool = True
+    
+    def setup(self):
+        if self.use_flash_attention:
+            attenBlock = EfficientAttention
+        else:
+            attenBlock = NormalAttention
+            
+        self.attention1 = attenBlock(
+         query_dim=self.query_dim,
+            heads=self.heads,
+            dim_head=self.dim_head,
+            name=f'Attention1',
+            precision=self.precision,
+            use_bias=self.use_bias,
+            dtype=self.dtype,
+            kernel_init=self.kernel_init,
+            force_fp32_for_softmax=self.force_fp32_for_softmax
+        )
+        self.attention2 = attenBlock(
+            query_dim=self.query_dim,
+            heads=self.heads,
+            dim_head=self.dim_head,
+            name=f'Attention2',
+            precision=self.precision,
+            use_bias=self.use_bias,
+            dtype=self.dtype,
+            kernel_init=self.kernel_init,
+            force_fp32_for_softmax=self.force_fp32_for_softmax
+        )
+        
+        self.ff = FlaxFeedForward(dim=self.query_dim)
+        self.norm1 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
+        self.norm2 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
+        self.norm3 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
+        
+    @nn.compact
+    def __call__(self, hidden_states, context=None):
+        if self.only_pure_attention:
+            return self.attention2(hidden_states, context)
+        
+        # self attention
+        if not self.use_cross_only:
+            hidden_states = hidden_states + self.attention1(self.norm1(hidden_states))
+        
+        # cross attention
+        hidden_states = hidden_states + self.attention2(self.norm2(hidden_states), context)
+        # feed forward
+        hidden_states = hidden_states + self.ff(self.norm3(hidden_states))
+        
+        return hidden_states
+
+class TransformerBlock(nn.Module):
+    heads: int = 4
+    dim_head: int = 32
+    use_linear_attention: bool = True
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    use_projection: bool = False
+    use_flash_attention:bool = False
+    use_self_and_cross:bool = True
+    only_pure_attention:bool = False
+    force_fp32_for_softmax: bool = True
+    kernel_init: Callable = kernel_init(1.0)
+    norm_inputs: bool = True
+    explicitly_add_residual: bool = True
+
+    @nn.compact
+    def __call__(self, x, context=None):
+        inner_dim = self.heads * self.dim_head
+        C = x.shape[-1]
+        if self.norm_inputs:
+            x = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)(x)
+        if self.use_projection == True:
+            if self.use_linear_attention:
+                projected_x = nn.Dense(features=inner_dim, 
+                                       use_bias=False, precision=self.precision, 
+                                       kernel_init=self.kernel_init,
+                                       dtype=self.dtype, name=f'project_in')(x)
+            else:
+                projected_x = nn.Conv(
+                    features=inner_dim, kernel_size=(1, 1),
+                    kernel_init=self.kernel_init,
+                    strides=(1, 1), padding='VALID', use_bias=False, dtype=self.dtype,
+                    precision=self.precision, name=f'project_in_conv',
+                )(x)
+        else:
+            projected_x = x
+            inner_dim = C
+            
+        context = projected_x if context is None else context
+
+        projected_x = BasicTransformerBlock(
+            query_dim=inner_dim,
+            heads=self.heads,
+            dim_head=self.dim_head,
+            name=f'Attention',
+            precision=self.precision,
+            use_bias=False,
+            dtype=self.dtype,
+            use_flash_attention=self.use_flash_attention,
+            use_cross_only=(not self.use_self_and_cross),
+            only_pure_attention=self.only_pure_attention,
+            force_fp32_for_softmax=self.force_fp32_for_softmax,
+            kernel_init=self.kernel_init
+        )(projected_x, context)
+        
+        if self.use_projection == True:
+            if self.use_linear_attention:
+                projected_x = nn.Dense(features=C, precision=self.precision, 
+                                       dtype=self.dtype, use_bias=False, 
+                                       kernel_init=self.kernel_init,
+                                       name=f'project_out')(projected_x)
+            else:
+                projected_x = nn.Conv(
+                    features=C, kernel_size=(1, 1),
+                    kernel_init=self.kernel_init,
+                    strides=(1, 1), padding='VALID', use_bias=False, dtype=self.dtype,
+                    precision=self.precision, name=f'project_out_conv',
+                )(projected_x)
+                
+        if self.only_pure_attention or self.explicitly_add_residual:
+            projected_x = x + projected_x
+            
+        out = projected_x
+        return out

flaxdiff/models/autoencoder/__init__.py

@@ -0,0 +1,2 @@
+from .autoencoder import AutoEncoder
+from .diffusers import StableDiffusionVAE

flaxdiff/models/autoencoder/autoencoder.py

@@ -0,0 +1,19 @@
+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, ConvLayer, Upsample, Downsample, PixelShuffle
+
+
+class AutoEncoder():
+    def encode(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
+        raise NotImplementedError
+    
+    def decode(self, z: jnp.ndarray, **kwargs) -> jnp.ndarray:
+        raise NotImplementedError
+    
+    def __call__(self, x: jnp.ndarray):
+        latents = self.encode(x)
+        reconstructions = self.decode(latents)
+        return reconstructions

flaxdiff/models/autoencoder/diffusers.py

@@ -0,0 +1,91 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from .autoencoder import AutoEncoder
+
+"""
+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.
+"""
+
+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
+        
+        pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
+            modelname,
+            revision=revision,
+            dtype=dtype,
+        )
+        
+        vae = pipeline.vae
+        
+        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,
+            double_z=True,
+            dtype=vae.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,
+        )
+        
+        quant_conv = nn.Conv(
+            2 * vae.config.latent_channels,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="VALID",
+            dtype=vae.dtype,
+        )
+
+        post_quant_conv = nn.Conv(
+            vae.config.latent_channels,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="VALID",
+            dtype=vae.dtype,
+        )
+        
+        self.enc = enc
+        self.dec = dec
+        self.post_quant_conv = post_quant_conv
+        self.quant_conv = quant_conv
+        self.params = params
+        self.scaling_factor = vae.scaling_factor
+        
+    def encode(self, images, rngkey: jax.random.PRNGKey = None):
+        latents = self.enc.apply({"params": self.params["vae"]['encoder']}, images, deterministic=True)
+        latents = self.quant_conv.apply({"params": self.params["vae"]['quant_conv']}, latents)
+        if rngkey is not None:
+            mean, log_std = jnp.split(latents, 2, axis=-1)
+            log_std = jnp.clip(log_std, -30, 20)
+            std = jnp.exp(0.5 * log_std)
+            latents = mean + std * jax.random.normal(rngkey, mean.shape, dtype=mean.dtype)
+            # print("Sampled")
+        else:
+            # return the mean
+            latents, _ = jnp.split(latents, 2, axis=-1)
+        latents *= self.scaling_factor
+        return latents
+    
+    def decode(self, latents):
+        latents = (1.0 / self.scaling_factor) * latents
+        latents = self.post_quant_conv.apply({"params": self.params["vae"]['post_quant_conv']}, latents)
+        return self.dec.apply({"params": self.params["vae"]['decoder']}, latents)

flaxdiff/models/autoencoder/simple_autoenc.py

@@ -0,0 +1,26 @@
+from typing import Any, List, Optional, Callable
+import jax
+import flax.linen as nn
+from jax import numpy as jnp
+from flax.typing import Dtype, PrecisionLike
+from .autoencoder import AutoEncoder
+
+class SimpleAutoEncoder(AutoEncoder):
+    latent_channels: int
+    feature_depths: List[int]=[64, 128, 256, 512]
+    attention_configs:list=[{"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}],
+    num_res_blocks: int=2
+    num_middle_res_blocks:int=1,
+    activation:Callable = jax.nn.swish
+    norm_groups:int=8
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    
+    # def encode(self, x: jnp.ndarray):
+        
+    
+    @nn.compact
+    def __call__(self, x: jnp.ndarray):
+        latents = self.encode(x)
+        reconstructions = self.decode(latents)
+        return reconstructions