flaxdiff 0.1.38__py3-none-any.whl → 0.2.0__py3-none-any.whl

flaxdiff/models/autoencoder/diffusers.py

@@ -22,7 +22,9 @@ class StableDiffusionVAE(AutoEncoder):
             dtype=dtype,
         )
         
-        # vae = pipeline.vae
+        self.modelname = modelname
+        self.revision = revision
+        self.dtype = dtype
         
         enc = FlaxEncoder(
             in_channels=vae.config.in_channels,
@@ -63,29 +65,90 @@ class StableDiffusionVAE(AutoEncoder):
             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
+        scaling_factor = vae.scaling_factor
+        print(f"Scaling factor: {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 encode_single_frame(images, rngkey: jax.random.PRNGKey = None):
+            latents = enc.apply({"params": params['encoder']}, images, deterministic=True)
+            latents = quant_conv.apply({"params": params['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)
+            else:
+                latents, _ = jnp.split(latents, 2, axis=-1)
+            latents *= scaling_factor
+            return latents
+        
+        def decode_single_frame(latents):
+            latents = (1.0 / scaling_factor) * latents
+            latents = post_quant_conv.apply({"params": params['post_quant_conv']}, latents)
+            return dec.apply({"params": params['decoder']}, latents)
+        
+        self.encode_single_frame = jax.jit(encode_single_frame)
+        self.decode_single_frame = jax.jit(decode_single_frame)
+        
+        # Calculate downscale factor by passing a dummy input through the encoder
+        print("Calculating downscale factor...")
+        dummy_input = jnp.ones((1, 128, 128, 3), dtype=dtype)
+        dummy_latents = self.encode_single_frame(dummy_input)
+        _, h, w, c = dummy_latents.shape
+        _, H, W, C = dummy_input.shape
+        self.__downscale_factor__ = H // h
+        self.__latent_channels__ = c
+        print(f"Downscale factor: {self.__downscale_factor__}")
+        print(f"Latent channels: {self.__latent_channels__}")
+
+    def __encode__(self, images, key: jax.random.PRNGKey = None, **kwargs):
+        """Encode a batch of images to latent representations.
+        
+        Implements the abstract method from the parent class.
+        
+        Args:
+            images: Image tensor of shape [B, H, W, C]
+            key: Optional random key for stochastic encoding
+            **kwargs: Additional arguments (unused)
+            
+        Returns:
+            Latent representations of shape [B, h, w, c]
+        """
+        return self.encode_single_frame(images, key)
+    
+    def __decode__(self, latents, **kwargs):
+        """Decode latent representations to images.
+        
+        Implements the abstract method from the parent class.
+        
+        Args:
+            latents: Latent tensor of shape [B, h, w, c]
+            **kwargs: Additional arguments (unused)
+            
+        Returns:
+            Decoded images of shape [B, H, W, C]
+        """
+        return self.decode_single_frame(latents)
+
+    @property
+    def downscale_factor(self) -> int:
+        """Returns the downscale factor for the encoder."""
+        return self.__downscale_factor__
+    
+    @property
+    def latent_channels(self) -> int:
+        """Returns the number of channels in the latent space."""
+        return self.__latent_channels__
+    
+    @property
+    def name(self) -> str:
+        """Get the name of the autoencoder model."""
+        return "stable_diffusion"
     
-    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)
+    def serialize(self):
+        """Serialize the model to a dictionary format."""
+        return {
+            "modelname": self.modelname,
+            "revision": self.revision,
+            "dtype": str(self.dtype),
+        }
+    

flaxdiff/models/autoencoder/simple_autoenc.py

@@ -6,21 +6,53 @@ from flax.typing import Dtype, PrecisionLike
 from .autoencoder import AutoEncoder
 
 class SimpleAutoEncoder(AutoEncoder):
+    """A simple autoencoder implementation using the abstract method pattern.
+    
+    This implementation allows for handling both image and video data through
+    the parent class's handling of video reshaping.
+    """
     latent_channels: int
     feature_depths: List[int]=[64, 128, 256, 512]
-    attention_configs:list=[{"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}],
+    attention_configs:list=[{"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}]
     num_res_blocks: int=2
-    num_middle_res_blocks:int=1,
+    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):
+    def __encode__(self, x: jnp.ndarray, **kwargs):
+        """Encode a batch of images to latent representations.
+        
+        Implements the abstract method from the parent class.
         
+        Args:
+            x: Image tensor of shape [B, H, W, C]
+            **kwargs: Additional arguments
+            
+        Returns:
+            Latent representations of shape [B, h, w, c]
+        """
+        # TODO: Implement the actual encoding logic for single frames
+        # This is just a placeholder implementation
+        B, H, W, C = x.shape
+        h, w = H // 8, W // 8  # Example downsampling factor
+        return jnp.zeros((B, h, w, self.latent_channels))
     
-    @nn.compact
-    def __call__(self, x: jnp.ndarray):
-        latents = self.encode(x)
-        reconstructions = self.decode(latents)
-        return reconstructions
+    def __decode__(self, z: jnp.ndarray, **kwargs):
+        """Decode latent representations to images.
+        
+        Implements the abstract method from the parent class.
+        
+        Args:
+            z: Latent tensor of shape [B, h, w, c]
+            **kwargs: Additional arguments
+            
+        Returns:
+            Decoded images of shape [B, H, W, C]
+        """
+        # TODO: Implement the actual decoding logic for single frames
+        # This is just a placeholder implementation
+        B, h, w, c = z.shape
+        H, W = h * 8, w * 8  # Example upsampling factor
+        return jnp.zeros((B, H, W, 3))

flaxdiff/models/common.py

@@ -108,13 +108,16 @@ class FourierEmbedding(nn.Module):
 class TimeProjection(nn.Module):
     features:int
     activation:Callable=jax.nn.gelu
-    kernel_init:Callable=kernel_init(1.0)
 
     @nn.compact
     def __call__(self, x):
-        x = nn.DenseGeneral(self.features, kernel_init=self.kernel_init)(x)
+        x = nn.DenseGeneral(
+            self.features, 
+        )(x)
         x = self.activation(x)
-        x = nn.DenseGeneral(self.features, kernel_init=self.kernel_init)(x)
+        x = nn.DenseGeneral(
+            self.features, 
+        )(x)
         x = self.activation(x)
         return x
 
@@ -123,7 +126,6 @@ class SeparableConv(nn.Module):
     kernel_size:tuple=(3, 3)
     strides:tuple=(1, 1)
     use_bias:bool=False
-    kernel_init:Callable=kernel_init(1.0)
     padding:str="SAME"
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
@@ -133,7 +135,7 @@ class SeparableConv(nn.Module):
         in_features = x.shape[-1]
         depthwise = nn.Conv(
             features=in_features, kernel_size=self.kernel_size,
-            strides=self.strides, kernel_init=self.kernel_init,
+            strides=self.strides,
             feature_group_count=in_features, use_bias=self.use_bias,
             padding=self.padding,
             dtype=self.dtype,
@@ -141,7 +143,7 @@ class SeparableConv(nn.Module):
         )(x)
         pointwise = nn.Conv(
             features=self.features, kernel_size=(1, 1),
-            strides=(1, 1), kernel_init=self.kernel_init,
+            strides=(1, 1),
             use_bias=self.use_bias,
             dtype=self.dtype,
             precision=self.precision
@@ -153,7 +155,6 @@ class ConvLayer(nn.Module):
     features:int
     kernel_size:tuple=(3, 3)
     strides:tuple=(1, 1)
-    kernel_init:Callable=kernel_init(1.0)
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
 
@@ -164,7 +165,6 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -183,7 +183,6 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -192,7 +191,6 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -206,7 +204,6 @@ class Upsample(nn.Module):
     activation:Callable=jax.nn.swish
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    kernel_init:Callable=kernel_init(1.0)
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -221,7 +218,6 @@ class Upsample(nn.Module):
             strides=(1, 1),
             dtype=self.dtype,
             precision=self.precision,
-            kernel_init=self.kernel_init
         )(out)
         if residual is not None:
             out = jnp.concatenate([out, residual], axis=-1)
@@ -233,7 +229,6 @@ class Downsample(nn.Module):
     activation:Callable=jax.nn.swish
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    kernel_init:Callable=kernel_init(1.0)
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -244,7 +239,6 @@ class Downsample(nn.Module):
             strides=(2, 2),
             dtype=self.dtype,
             precision=self.precision,
-            kernel_init=self.kernel_init
         )(x)
         if residual is not None:
             if residual.shape[1] > out.shape[1]:
@@ -269,7 +263,6 @@ class ResidualBlock(nn.Module):
     direction:str=None
     res:int=2
     norm_groups:int=8
-    kernel_init:Callable=kernel_init(1.0)
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     named_norms:bool=False
@@ -296,7 +289,6 @@ class ResidualBlock(nn.Module):
             features=self.features,
             kernel_size=self.kernel_size,
             strides=self.strides,
-            kernel_init=self.kernel_init,
             name="conv1",
             dtype=self.dtype,
             precision=self.precision
@@ -321,7 +313,6 @@ class ResidualBlock(nn.Module):
             features=self.features,
             kernel_size=self.kernel_size,
             strides=self.strides,
-            kernel_init=self.kernel_init,
             name="conv2",
             dtype=self.dtype,
             precision=self.precision
@@ -333,7 +324,6 @@ class ResidualBlock(nn.Module):
                 features=self.features,
                 kernel_size=(1, 1),
                 strides=1,
-                kernel_init=self.kernel_init,
                 name="residual_conv",
                 dtype=self.dtype,
                 precision=self.precision

flaxdiff/models/simple_unet.py

@@ -10,17 +10,16 @@ from functools import partial
 
 class Unet(nn.Module):
     output_channels:int=3
-    emb_features:int=64*4,
-    feature_depths:list=[64, 128, 256, 512],
-    attention_configs:list=[{"heads":8}, {"heads":8}, {"heads":8}, {"heads":8}],
-    num_res_blocks:int=2,
-    num_middle_res_blocks:int=1,
+    emb_features:int=64*4
+    feature_depths:list=(64, 128, 256, 512)
+    attention_configs:list=({"heads":8}, {"heads":8}, {"heads":8}, {"heads":8})
+    num_res_blocks:int=2
+    num_middle_res_blocks:int=1
     activation:Callable = jax.nn.swish
     norm_groups:int=8
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     named_norms: bool = False # This is for backward compatibility reasons; older checkpoints have named norms
-    kernel_init: Callable = partial(kernel_init, dtype=jnp.float32)
 
     def setup(self):
         if self.norm_groups > 0:
@@ -50,7 +49,6 @@ class Unet(nn.Module):
             features=self.feature_depths[0],
             kernel_size=(3, 3),
             strides=(1, 1),
-            kernel_init=self.kernel_init(scale=1.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)
@@ -65,7 +63,6 @@ class Unet(nn.Module):
                     down_conv_type,
                     name=f"down_{i}_residual_{j}",
                     features=dim_in,
-                    kernel_init=self.kernel_init(scale=1.0),
                     kernel_size=(3, 3),
                     strides=(1, 1),
                     activation=self.activation,
@@ -85,7 +82,6 @@ class Unet(nn.Module):
                                         force_fp32_for_softmax=attention_config.get("force_fp32_for_softmax", False),
                                         norm_inputs=attention_config.get("norm_inputs", True),
                                         explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
-                                        kernel_init=self.kernel_init(scale=1.0),
                                         name=f"down_{i}_attention_{j}")(x, textcontext)
                 # print("down residual for feature level", i, "is of shape", x.shape, "features", dim_in)
                 downs.append(x)
@@ -108,7 +104,6 @@ class Unet(nn.Module):
                 middle_conv_type,
                 name=f"middle_res1_{j}",
                 features=middle_dim_out,
-                kernel_init=self.kernel_init(scale=1.0),
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 activation=self.activation,
@@ -129,13 +124,11 @@ class Unet(nn.Module):
                                     force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
                                     norm_inputs=middle_attention.get("norm_inputs", True),
                                     explicitly_add_residual=middle_attention.get("explicitly_add_residual", True),
-                                    kernel_init=self.kernel_init(scale=1.0),
                                     name=f"middle_attention_{j}")(x, textcontext)
             x = ResidualBlock(
                 middle_conv_type,
                 name=f"middle_res2_{j}",
                 features=middle_dim_out,
-                kernel_init=self.kernel_init(scale=1.0),
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 activation=self.activation,
@@ -157,7 +150,6 @@ class Unet(nn.Module):
                     up_conv_type,# if j == 0 else "separable",
                     name=f"up_{i}_residual_{j}",
                     features=dim_out,
-                    kernel_init=self.kernel_init(scale=1.0),
                     kernel_size=kernel_size,
                     strides=(1, 1),
                     activation=self.activation,
@@ -177,7 +169,6 @@ class Unet(nn.Module):
                                         force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
                                         norm_inputs=attention_config.get("norm_inputs", True),
                                         explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
-                                        kernel_init=self.kernel_init(scale=1.0),
                                         name=f"up_{i}_attention_{j}")(x, textcontext)
             # print("Upscaling ", i, x.shape)
             if i != len(feature_depths) - 1:
@@ -196,7 +187,6 @@ class Unet(nn.Module):
             features=self.feature_depths[0],
             kernel_size=(3, 3),
             strides=(1, 1),
-            kernel_init=self.kernel_init(scale=1.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)
@@ -207,7 +197,6 @@ class Unet(nn.Module):
             conv_type,
             name="final_residual",
             features=self.feature_depths[0],
-            kernel_init=self.kernel_init(scale=1.0),
             kernel_size=(3,3),
             strides=(1, 1),
             activation=self.activation,
@@ -226,7 +215,7 @@ class Unet(nn.Module):
             kernel_size=(3, 3),
             strides=(1, 1),
             # activation=jax.nn.mish
-            kernel_init=self.kernel_init(scale=0.0),
+            # kernel_init=self.kernel_init(scale=0.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)

flaxdiff/models/simple_vit.py

@@ -23,7 +23,6 @@ class PatchEmbedding(nn.Module):
     embedding_dim: int
     dtype: Any = jnp.float32
     precision: Any = jax.lax.Precision.HIGH
-    kernel_init: Callable = partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x):
@@ -34,7 +33,6 @@ class PatchEmbedding(nn.Module):
                     kernel_size=(self.patch_size, self.patch_size), 
                     strides=(self.patch_size, self.patch_size),
                     dtype=self.dtype,
-                    kernel_init=self.kernel_init(),
                     precision=self.precision)(x)
         x = jnp.reshape(x, (batch, -1, self.embedding_dim))
         return x
@@ -53,7 +51,7 @@ class PositionalEncoding(nn.Module):
 class UViT(nn.Module):
     output_channels:int=3
     patch_size: int = 16
-    emb_features:int=768,
+    emb_features:int=768
     num_layers: int = 12
     num_heads: int = 12
     dropout_rate: float = 0.1
@@ -67,7 +65,7 @@ class UViT(nn.Module):
     norm_groups:int=8
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    kernel_init: Callable = partial(kernel_init, scale=1.0)
+    # kernel_init: Callable = partial(kernel_init, scale=1.0)
     add_residualblock_output: bool = False
     norm_inputs: bool = False
     explicitly_add_residual: bool = True
@@ -88,10 +86,10 @@ class UViT(nn.Module):
 
         # Patch embedding
         x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
-                           dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init)(x)
+                           dtype=self.dtype, precision=self.precision)(x)
         num_patches = x.shape[1]
         
-        context_emb = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
+        context_emb = nn.DenseGeneral(features=self.emb_features, 
                                dtype=self.dtype, precision=self.precision)(textcontext)
         num_text_tokens = textcontext.shape[1]
         
@@ -116,7 +114,7 @@ class UViT(nn.Module):
                                  only_pure_attention=False,
                                  norm_inputs=self.norm_inputs,
                                  explicitly_add_residual=self.explicitly_add_residual,
-                                 kernel_init=self.kernel_init())(x)
+                                 )(x)
             skips.append(x)
             
         # Middle block
@@ -126,12 +124,12 @@ class UViT(nn.Module):
                              only_pure_attention=False,
                             norm_inputs=self.norm_inputs,
                             explicitly_add_residual=self.explicitly_add_residual,
-                             kernel_init=self.kernel_init())(x)
+                            )(x)
         
         # # Out blocks
         for i in range(self.num_layers // 2):
             x = jnp.concatenate([x, skips.pop()], axis=-1)
-            x = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
+            x = nn.DenseGeneral(features=self.emb_features, 
                                    dtype=self.dtype, precision=self.precision)(x)
             x = TransformerBlock(heads=self.num_heads, dim_head=self.emb_features // self.num_heads, 
                                  dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, 
@@ -139,13 +137,13 @@ class UViT(nn.Module):
                                  only_pure_attention=False,
                                  norm_inputs=self.norm_inputs,
                                  explicitly_add_residual=self.explicitly_add_residual,
-                                 kernel_init=self.kernel_init())(x)
+                                 )(x)
         
         # print(f'Shape of x after transformer blocks: {x.shape}')
         x = self.norm()(x)
         
         patch_dim = self.patch_size ** 2 * self.output_channels
-        x = nn.Dense(features=patch_dim, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
+        x = nn.Dense(features=patch_dim, dtype=self.dtype, precision=self.precision)(x)
         x = x[:, 1 + num_text_tokens:, :]
         x = unpatchify(x, channels=self.output_channels)
         
@@ -159,7 +157,6 @@ class UViT(nn.Module):
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 # activation=jax.nn.mish
-                kernel_init=self.kernel_init(scale=0.0),
                 dtype=self.dtype,
                 precision=self.precision
             )(x)
@@ -173,7 +170,6 @@ class UViT(nn.Module):
             kernel_size=(3, 3),
             strides=(1, 1),
             # activation=jax.nn.mish
-            kernel_init=self.kernel_init(scale=0.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)