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

flaxdiff/models/attention.py

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

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

flaxdiff/models/autoencoder/autoencoder.py

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

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

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