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

flaxdiff/inputs/__init__.py

@@ -0,0 +1,173 @@
+import jax
+import jax.numpy as jnp
+import flax.struct as struct
+import flax.linen as nn
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from dataclasses import dataclass
+from functools import partial
+import numpy as np
+from jax.sharding import Mesh, PartitionSpec as P
+from abc import ABC, abstractmethod
+
+from flaxdiff.models.autoencoder import AutoEncoder
+from .encoders import *
+
+@dataclass
+class ConditionalInputConfig:
+    """Class representing a conditional input for the model."""
+    encoder: ConditioningEncoder
+    conditioning_data_key: str = None       # Key in the batch for this conditioning input
+    pretokenized: bool = False
+    unconditional_input: Any = None
+    model_key_override: Optional[str] = None  # Optional key override for the model
+    
+    __uncond_cache__ = None  # Cache for unconditional input
+    
+    def __post_init__(self):
+        if self.unconditional_input is not None:
+            uncond = self.encoder([self.unconditional_input])
+        else:
+            uncond = self.encoder([""])  # Default empty text
+        self.__uncond_cache__ = uncond  # Cache the unconditional input
+    
+    def __call__(self, batch_data):
+        """Process batch data to produce conditioning."""
+        key =  self.conditioning_data_key if self.conditioning_data_key else self.encoder.key
+        if self.pretokenized:
+            return self.encoder.encode_from_tokens(batch_data[key])
+        return self.encoder(batch_data[key])
+    
+    def get_unconditional(self):
+        """Get unconditional version of this input."""
+        return self.__uncond_cache__
+    
+    def serialize(self):
+        """Serialize the configuration."""
+        serialized_config = {
+            "encoder": self.encoder.serialize(),
+            "encoder_key": self.encoder.key,
+            "conditioning_data_key": self.conditioning_data_key,
+            "unconditional_input": self.unconditional_input,
+            "model_key_override": self.model_key_override,
+        }
+        return serialized_config
+    
+    @staticmethod
+    def deserialize(serialized_config):
+        """Deserialize the configuration."""
+        encoder_key = serialized_config["encoder_key"]
+        encoder_class = CONDITIONAL_ENCODERS_REGISTRY.get(encoder_key)
+        if encoder_class is None:
+            raise ValueError(f"Unknown encoder type: {encoder_key}")
+        
+        # Create the encoder instance
+        encoder = encoder_class.deserialize(serialized_config["encoder"])
+        # Deserialize the rest of the configuration
+        conditioning_data_key = serialized_config.get("conditioning_data_key")
+        unconditional_input = serialized_config.get("unconditional_input")
+        model_key_override = serialized_config.get("model_key_override")
+        return ConditionalInputConfig(
+            encoder=encoder,
+            conditioning_data_key=conditioning_data_key,
+            unconditional_input=unconditional_input,
+            model_key_override=model_key_override,
+        )
+    
+@dataclass
+class DiffusionInputConfig:
+    """Configuration for the input data."""
+    sample_data_key: str         # Key in the batch for the sample data
+    sample_data_shape: Tuple[int, ...]
+    conditions: List[ConditionalInputConfig]
+    
+    def get_input_shapes(
+        self, 
+        autoencoder: AutoEncoder = None, 
+        sample_model_key:str = 'x',
+        time_embeddings_model_key:str = 'temb',
+    ) -> Dict[str, Tuple[int, ...]]:
+        """Get the shapes of the input data."""
+        if len(self.sample_data_shape) == 3:
+            H, W, C = self.sample_data_shape
+        elif len(self.sample_data_shape) == 4:
+            T, H, W, C = self.sample_data_shape
+        else:
+            raise ValueError(f"Unsupported shape for sample data {self.sample_data_shape}")
+        if autoencoder is not None:
+            downscale_factor = autoencoder.downscale_factor
+            H = H // downscale_factor
+            W = W // downscale_factor
+            C = autoencoder.latent_channels
+        
+        input_shapes = {
+            sample_model_key: (H, W, C),
+            time_embeddings_model_key: (),
+        }
+        for cond in self.conditions:
+            # Get the shape of the conditioning data by calling the get_unconditional method
+            unconditional = cond.get_unconditional()
+            key = cond.model_key_override if cond.model_key_override else cond.encoder.key
+            input_shapes[key] = unconditional[0].shape
+            
+        print(f"Calculated input shapes: {input_shapes}")
+        return input_shapes
+    
+    def get_unconditionals(self):
+        """Get unconditional inputs for all conditions."""
+        unconditionals = []
+        for cond in self.conditions:
+            uncond = cond.get_unconditional()
+            unconditionals.append(uncond)
+        return unconditionals
+    
+    def process_conditioning(self, batch_data, uncond_mask: Optional[jnp.ndarray] = None):
+        """Process the conditioning data."""
+        results = []
+            
+        for cond in self.conditions:
+            cond_embeddings = cond(batch_data)
+            if uncond_mask is not None:
+                assert len(uncond_mask) == len(cond_embeddings), "Unconditional mask length must match the batch size."
+                uncond_embedding = cond.get_unconditional()
+                    
+                # Reshape uncond_mask to be broadcastable with the conditioning embeddings
+                # If cond_embeddings has shape (B, T, D), reshape uncond_mask to (B, 1, 1)
+                broadcast_shape = [len(uncond_mask)] + [1] * (cond_embeddings.ndim - 1)
+                reshaped_mask = jnp.reshape(uncond_mask, broadcast_shape)
+                    
+                # Repeat uncond_embedding to match batch size
+                batch_size = len(cond_embeddings)
+                repeated_uncond = jnp.repeat(uncond_embedding, batch_size, axis=0)
+                    
+                # Apply unconditional embedding based on the mask
+                cond_embeddings = jnp.where(reshaped_mask, repeated_uncond, cond_embeddings)
+                
+            results.append(cond_embeddings)
+        return results
+    
+    def serialize(self):
+        """Serialize the configuration."""
+        serialized_config = {
+            "sample_data_key": self.sample_data_key,
+            "sample_data_shape": self.sample_data_shape,
+            "conditions": [cond.serialize() for cond in self.conditions],
+        }
+        return serialized_config
+    
+    @staticmethod
+    def deserialize(serialized_config):
+        """Deserialize the configuration."""
+        sample_data_key = serialized_config["sample_data_key"]
+        sample_data_shape = tuple(serialized_config["sample_data_shape"])
+        conditions = serialized_config["conditions"]
+        
+        # Deserialize each condition
+        deserialized_conditions = []
+        for cond in conditions:
+            deserialized_conditions.append(ConditionalInputConfig.deserialize(cond))
+        
+        return DiffusionInputConfig(
+            sample_data_key=sample_data_key,
+            sample_data_shape=sample_data_shape,
+            conditions=deserialized_conditions,
+        )

flaxdiff/inputs/encoders.py

@@ -0,0 +1,98 @@
+import jax.numpy as jnp
+import flax.linen as nn
+from typing import Callable
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+
+@dataclass
+class ConditioningEncoder(ABC):
+    model: nn.Module
+    tokenizer: Callable
+    
+    @property
+    def key(self):
+        name = self.tokenizer.__name__
+        # Remove the 'Encoder' suffix from the name and lowercase it
+        if name.endswith("Encoder"):
+            name = name[:-7].lower()
+        return name
+
+    def __call__(self, data):
+        tokens = self.tokenize(data)
+        outputs = self.encode_from_tokens(tokens)
+        return outputs
+        
+    def encode_from_tokens(self, tokens):
+        outputs = self.model(input_ids=tokens['input_ids'],
+                        attention_mask=tokens['attention_mask'])
+        last_hidden_state = outputs.last_hidden_state
+        return last_hidden_state
+    
+    def tokenize(self, data):
+        tokens = self.tokenizer(data, padding="max_length",
+                        max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="np")
+        return tokens
+    
+    @abstractmethod
+    def serialize(self):
+        """Serialize the encoder configuration."""
+        pass
+    
+    @staticmethod
+    @abstractmethod
+    def deserialize(serialized_config):
+        """Deserialize the encoder configuration."""
+        pass
+    
+@dataclass
+class TextEncoder(ConditioningEncoder):
+    """Text Encoder."""
+    @property
+    def key(self):
+        return "text"
+    
+@dataclass
+class CLIPTextEncoder(TextEncoder):
+    """CLIP Text Encoder."""
+    modelname: str
+    backend: str
+    
+    @staticmethod
+    def from_modelname(modelname: str = "openai/clip-vit-large-patch14", backend: str="jax"):
+        from transformers import (
+            CLIPTextModel,
+            FlaxCLIPTextModel,
+            AutoTokenizer,
+        )
+        modelname = "openai/clip-vit-large-patch14"
+        if backend == "jax":
+            model = FlaxCLIPTextModel.from_pretrained(
+                modelname, dtype=jnp.bfloat16)
+        else:
+            model = CLIPTextModel.from_pretrained(modelname)
+        tokenizer = AutoTokenizer.from_pretrained(modelname, dtype=jnp.float16)
+        return CLIPTextEncoder(
+            model=model,
+            tokenizer=tokenizer,
+            modelname=modelname,
+            backend=backend
+        )
+    
+    def serialize(self):
+        """Serialize the encoder configuration."""
+        serialized_config = {
+            "modelname": self.modelname,
+            "backend": self.backend,
+        }
+        return serialized_config
+    
+    @staticmethod
+    def deserialize(serialized_config):
+        """Deserialize the encoder configuration."""
+        modelname = serialized_config["modelname"]
+        backend = serialized_config["backend"]
+        return CLIPTextEncoder.from_modelname(modelname=modelname, backend=backend)
+    
+CONDITIONAL_ENCODERS_REGISTRY = {
+    "text": CLIPTextEncoder,
+}

flaxdiff/models/__init__.py

@@ -1 +1,2 @@
-from .simple_unet import *
+from .simple_unet import *
+# from .video_unet import FlaxUNet3DConditionModel, BCHWModelWrapper, FlaxTemporalConvLayer

flaxdiff/models/autoencoder/autoencoder.py

@@ -1,19 +1,151 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from typing import Dict, Callable, Sequence, Any, Union
+from typing import Dict, Callable, Sequence, Any, Union, Optional
 import einops
 from ..common import kernel_init, ConvLayer, Upsample, Downsample, PixelShuffle
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
 
-
-class AutoEncoder():
-    def encode(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
+@dataclass
+class AutoEncoder(ABC):
+    """Base class for autoencoder models with video support.
+    
+    This class defines the interface for autoencoders and provides
+    video handling functionality, allowing child classes to focus
+    on implementing the core encoding/decoding for individual frames.
+    """
+    @abstractmethod
+    def __encode__(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
+        """Abstract method for encoding a batch of images.
+        
+        Child classes must implement this method to perform the actual encoding.
+        
+        Args:
+            x: Input tensor of shape [B, H, W, C] (batch of images)
+            **kwargs: Additional arguments for the encoding process
+            
+        Returns:
+            Encoded latent representation
+        """
+        raise NotImplementedError
+    
+    @abstractmethod
+    def __decode__(self, z: jnp.ndarray, **kwargs) -> jnp.ndarray:
+        """Abstract method for decoding a batch of latents.
+        
+        Child classes must implement this method to perform the actual decoding.
+        
+        Args:
+            z: Latent tensor of shape [B, h, w, c] (encoded representation)
+            **kwargs: Additional arguments for the decoding process
+            
+        Returns:
+            Decoded images
+        """
         raise NotImplementedError
     
-    def decode(self, z: jnp.ndarray, **kwargs) -> jnp.ndarray:
+    def encode(self, x: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs) -> jnp.ndarray:
+        """Encode input data, with special handling for video data.
+        
+        This method handles both standard image batches and video data (5D tensors).
+        For videos, it reshapes the input, processes each frame, and then restores
+        the temporal dimension.
+        
+        Args:
+            x: Input tensor, either [B, H, W, C] for images or [B, T, H, W, C] for videos
+            key: Optional random key for stochastic encoding
+            **kwargs: Additional arguments passed to __encode__
+            
+        Returns:
+            Encoded representation with the same batch and temporal dimensions as input
+        """
+        # Check for video data (5D tensor)
+        is_video = len(x.shape) == 5
+        
+        if is_video:
+            # Extract dimensions for reshaping
+            batch_size, seq_len, height, width, channels = x.shape
+            
+            # Reshape to [B*T, H, W, C] to process as regular images
+            x_reshaped = x.reshape(-1, height, width, channels)
+            
+            # Encode all frames
+            latent = self.__encode__(x_reshaped, key=key, **kwargs)
+            
+            # Reshape back to include temporal dimension [B, T, h, w, c]
+            latent_shape = latent.shape
+            return latent.reshape(batch_size, seq_len, *latent_shape[1:])
+        else:
+            # Standard image processing
+            return self.__encode__(x, key=key, **kwargs)
+    
+    def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs) -> jnp.ndarray:
+        """Decode latent representations, with special handling for video data.
+        
+        This method handles both standard image latents and video latents (5D tensors).
+        For videos, it reshapes the input, processes each frame, and then restores
+        the temporal dimension.
+        
+        Args:
+            z: Latent tensor, either [B, h, w, c] for images or [B, T, h, w, c] for videos
+            key: Optional random key for stochastic decoding
+            **kwargs: Additional arguments passed to __decode__
+            
+        Returns:
+            Decoded output with the same batch and temporal dimensions as input
+        """
+        # Check for video data (5D tensor)
+        is_video = len(z.shape) == 5
+        
+        if is_video:
+            # Extract dimensions for reshaping
+            batch_size, seq_len, height, width, channels = z.shape
+            
+            # Reshape to [B*T, h, w, c] to process as regular latents
+            z_reshaped = z.reshape(-1, height, width, channels)
+            
+            # Decode all frames
+            decoded = self.__decode__(z_reshaped, key=key, **kwargs)
+            
+            # Reshape back to include temporal dimension [B, T, H, W, C]
+            decoded_shape = decoded.shape
+            return decoded.reshape(batch_size, seq_len, *decoded_shape[1:])
+        else:
+            # Standard latent processing
+            return self.__decode__(z, key=key, **kwargs)
+    
+    def __call__(self, x: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs):
+        """Encode and then decode the input (autoencoder).
+        
+        Args:
+            x: Input tensor, either [B, H, W, C] for images or [B, T, H, W, C] for videos
+            key: Optional random key for stochastic encoding/decoding
+            **kwargs: Additional arguments for encoding and decoding
+            
+        Returns:
+            Reconstructed output with the same dimensions as input
+        """
+        if key is not None:
+            encode_key, decode_key = jax.random.split(key)
+        else:
+            encode_key = decode_key = None
+            
+        # Encode then decode
+        z = self.encode(x, key=encode_key, **kwargs)
+        return self.decode(z, key=decode_key, **kwargs)
+    
+    @property
+    def spatial_scale(self) -> int:
+        """Get the spatial scale factor between input and latent spaces."""
+        return getattr(self, "_spatial_scale", None)
+    
+    @property
+    def name(self) -> str:
+        """Get the name of the autoencoder model."""
         raise NotImplementedError
     
-    def __call__(self, x: jnp.ndarray):
-        latents = self.encode(x)
-        reconstructions = self.decode(latents)
-        return reconstructions
+    @abstractmethod
+    def serialize(self) -> Dict[str, Any]:
+        """Serialize the model parameters and configuration."""
+        raise NotImplementedError

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/simple_unet.py

@@ -10,11 +10,11 @@ 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

flaxdiff/models/simple_vit.py

@@ -51,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