optical_embeddings/vision/

clip.rs

//! CLIP vision encoder used by DeepSeek-OCR.
//!
//! This module implements the CLIP (Contrastive Language-Image Pre-training) vision encoder
//! component of the DeepSeek-OCR pipeline. CLIP provides global attention processing for
//! semantic understanding of visual features.
//!
//! # Architecture
//!
//! The CLIP encoder consists of:
//! - **Vision embeddings**: Patch embedding with class token and position embeddings
//! - **Pre-LayerNorm**: Applied before transformer blocks
//! - **Transformer blocks**: Stack of self-attention layers with Quick GELU activation
//!
//! # Configuration
//!
//! Supports multiple model sizes:
//! - **CLIP-Base**: 768-dim, 12 layers, 12 heads
//! - **CLIP-Large**: 1024-dim, 24 layers, 16 heads (default)
//!
//! # Example
//!
//! ```ignore
//! use optical_embeddings::vision::ClipEncoder;
//! use optical_embeddings::config::ClipConfig;
//!
//! let config = ClipConfig::large();
//! let encoder = ClipEncoder::new(&config, &device);
//! let features = encoder.forward(image_tensor, None);
//! ```

use burn::nn::{
    conv::{Conv2d, Conv2dConfig},
    Embedding, EmbeddingConfig, LayerNorm, LayerNormConfig,
};
use burn::prelude::*;
use burn::tensor::{backend::Backend, Distribution, Tensor};
use log::{debug, info, trace};

use super::attention::Attention;
use crate::config::ClipConfig;

/// CLIP vision embeddings layer.
///
/// Converts input images to patch embeddings with class token and position encodings.
/// Implements the embedding strategy used in Vision Transformers (ViT).
///
/// # Components
///
/// - `class_embedding`: Learnable class token prepended to patch sequence
/// - `patch_embedding`: Convolutional layer that splits image into patches
/// - `position_embedding`: Learnable positional encodings for each patch + class token
#[derive(Module, Debug)]
struct ClipVisionEmbeddings<B: Backend> {
    class_embedding: Tensor<B, 1>,
    patch_embedding: Conv2d<B>,
    position_embedding: Embedding<B>,
    embed_dim: usize,
}

impl<B: Backend> ClipVisionEmbeddings<B> {
    /// Creates a new CLIP vision embeddings layer.
    ///
    /// # Arguments
    ///
    /// * `cfg` - CLIP configuration specifying embedding dimensions and patch size
    /// * `device` - Device to initialize tensors on (CPU/GPU)
    ///
    /// # Returns
    ///
    /// Initialized embeddings layer ready for forward pass
    fn new(cfg: &ClipConfig, device: &B::Device) -> Self {
        info!(
            "Creating CLIP vision embeddings: hidden_size={}, image_size={}, patch_size={}",
            cfg.hidden_size, cfg.image_size, cfg.patch_size
        );

        let embed_dim = cfg.hidden_size;
        let cls = Tensor::random([embed_dim], Distribution::Normal(0.0, 0.02), device);
        let pe = Conv2dConfig::new([3, embed_dim], [cfg.patch_size, cfg.patch_size])
            .with_stride([cfg.patch_size, cfg.patch_size])
            .with_bias(false)
            .init(device);
        let num_patches = (cfg.image_size / cfg.patch_size).pow(2);

        debug!(
            "CLIP embeddings: num_patches={}, embed_dim={}",
            num_patches, embed_dim
        );

        let pos = EmbeddingConfig::new(num_patches + 1, embed_dim).init(device);
        Self {
            class_embedding: cls,
            patch_embedding: pe,
            position_embedding: pos,
            embed_dim,
        }
    }

    /// Forward pass through embeddings layer.
    ///
    /// # Arguments
    ///
    /// * `pixels` - Input image tensor of shape `[B, 3, H, W]`
    /// * `patch_embeds` - Optional pre-computed patch embeddings (currently unused)
    ///
    /// # Returns
    ///
    /// Embedded sequence of shape `[B, num_patches+1, embed_dim]` where the first
    /// token is the class token and remaining tokens are patch embeddings with
    /// position encodings added.
    fn forward(&self, pixels: Tensor<B, 4>, patch_embeds: Option<Tensor<B, 4>>) -> Tensor<B, 3> {
        trace!("CLIP embeddings forward: input shape={:?}", pixels.dims());

        let _ = pixels.dims()[0];

        // Get patch embeddings: either provided or compute from pixels
        let patches = if let Some(_) = patch_embeds {
            debug!("Using provided patch embeddings");
            // If SAM features provided, they need to be resized/projected to match embed_dim
            // For now, we'll compute fresh patch embeddings
            self.patch_embedding.forward(pixels)
        } else {
            trace!("Computing fresh patch embeddings");
            self.patch_embedding.forward(pixels)
        };

        // patches: [B, embed_dim, H_patches, W_patches]
        let shp = patches.dims();
        let batch = shp[0];
        let embed_dim = shp[1];
        let h_patches = shp[2];
        let w_patches = shp[3];
        let num_patches = h_patches * w_patches;

        debug!("Patch embedding output: batch={}, embed_dim={}, h_patches={}, w_patches={}, total_patches={}", 
               batch, embed_dim, h_patches, w_patches, num_patches);

        // Reshape to [B, embed_dim, num_patches] then transpose to [B, num_patches, embed_dim]
        let patches = patches
            .reshape([batch, embed_dim, num_patches])
            .swap_dims(1, 2);

        // Prepare class token: [embed_dim] -> [1, 1, embed_dim] -> [B, 1, embed_dim]
        let cls = self
            .class_embedding
            .clone()
            .reshape([1, 1, self.embed_dim])
            .repeat(&[batch]);

        // Concatenate class token with patches: [B, num_patches+1, embed_dim]
        let embs = Tensor::cat(vec![cls, patches], 1);
        let seq_len = embs.dims()[1];

        debug!(
            "After concatenation: seq_len={} (1 cls + {} patches)",
            seq_len, num_patches
        );

        // Create position IDs: [0, 1, 2, ..., seq_len-1]
        let pos_ids = Tensor::arange(0..(seq_len as i64), &embs.device())
            .reshape([1, seq_len])
            .repeat(&[batch]);

        // Add position embeddings
        let result = embs + self.position_embedding.forward(pos_ids);
        trace!("CLIP embeddings output shape: {:?}", result.dims());
        result
    }
}

/// CLIP transformer block implementing pre-LayerNorm architecture.
///
/// Each block consists of:
/// 1. LayerNorm → Multi-head self-attention → Residual connection
/// 2. LayerNorm → Feed-forward network (with Quick GELU) → Residual connection
///
/// # Quick GELU
///
/// Uses the approximation: `x * sigmoid(1.702 * x)` instead of standard GELU
/// for computational efficiency as specified in CLIP.
#[derive(Module, Debug)]
struct ClipTransformerBlock<B: Backend> {
    ln1: LayerNorm<B>,
    attn: Attention<B>,
    ln2: LayerNorm<B>,
    ff1: burn::nn::Linear<B>,
    ff2: burn::nn::Linear<B>,
}

impl<B: Backend> ClipTransformerBlock<B> {
    /// Creates a new CLIP transformer block.
    ///
    /// # Arguments
    ///
    /// * `cfg` - CLIP configuration with layer dimensions and hyperparameters
    /// * `device` - Device for tensor initialization
    fn new(cfg: &ClipConfig, device: &B::Device) -> Self {
        trace!(
            "Creating CLIP transformer block: hidden_size={}, num_heads={}, ffn_hidden={}",
            cfg.hidden_size,
            cfg.num_attention_heads,
            cfg.ffn_hidden_size
        );

        let ln1 = LayerNormConfig::new(cfg.hidden_size)
            .with_epsilon(cfg.layernorm_epsilon as f64)
            .init(device);
        let ln2 = LayerNormConfig::new(cfg.hidden_size)
            .with_epsilon(cfg.layernorm_epsilon as f64)
            .init(device);
        let attn = Attention::new(cfg.hidden_size, cfg.num_attention_heads, device);
        let ff1 = burn::nn::LinearConfig::new(cfg.hidden_size, cfg.ffn_hidden_size)
            .with_bias(true)
            .init(device);
        let ff2 = burn::nn::LinearConfig::new(cfg.ffn_hidden_size, cfg.hidden_size)
            .with_bias(true)
            .init(device);
        Self {
            ln1,
            attn,
            ln2,
            ff1,
            ff2,
        }
    }

    /// Quick GELU activation function.
    ///
    /// Approximation of GELU using: `x * sigmoid(1.702 * x)`
    ///
    /// # Arguments
    ///
    /// * `x` - Input tensor
    ///
    /// # Returns
    ///
    /// Activated tensor with same shape as input
    fn quick_gelu(x: Tensor<B, 3>) -> Tensor<B, 3> {
        let s = burn::tensor::activation::sigmoid(x.clone() * 1.702);
        x * s
    }

    /// Forward pass through transformer block.
    ///
    /// Applies pre-LayerNorm architecture:
    /// 1. x = x + Attention(LayerNorm(x))
    /// 2. x = x + FFN(LayerNorm(x))
    ///
    /// # Arguments
    ///
    /// * `x` - Input tensor of shape `[B, seq_len, hidden_size]`
    ///
    /// # Returns
    ///
    /// Output tensor of same shape as input
    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        trace!("CLIP transformer block forward: input shape={:?}", x.dims());
        let h = x.clone() + self.attn.forward(self.ln1.forward(x.clone()));
        let y = self.ff2.forward(Self::quick_gelu(
            self.ff1.forward(self.ln2.forward(h.clone())),
        ));
        h + y
    }
}

/// CLIP vision encoder for DeepSeek-OCR.
///
/// Implements the CLIP-Large architecture (or configurable variants) for global
/// semantic understanding of visual features. Processes patch embeddings through
/// a stack of transformer blocks with pre-LayerNorm and Quick GELU activation.
///
/// # Architecture Details
///
/// - **Input**: RGB images `[B, 3, H, W]`
/// - **Patch embedding**: Splits image into non-overlapping patches
/// - **Class token**: Learnable token prepended to sequence
/// - **Position encoding**: Added to all tokens
/// - **Transformer**: Stack of attention + FFN blocks
/// - **Output**: Sequence of embeddings `[B, num_patches+1, hidden_size]`
///
/// # Default Configuration (CLIP-Large)
///
/// - 1024-dimensional embeddings
/// - 24 transformer layers
/// - 16 attention heads per layer
/// - 4096-dimensional feed-forward hidden layer
/// - 14×14 patch size on 224×224 images (256 patches)
///
/// # Example
///
/// ```ignore
/// let config = ClipConfig::large();
/// let encoder = ClipEncoder::new(&config, &device);
///
/// let image = Tensor::zeros(, &device);[1]
/// let features = encoder.forward(image, None);
/// // features shape:  (256 patches + 1 class token)[1]
/// ```
#[derive(Module, Debug)]
pub struct ClipEncoder<B: Backend> {
    emb: ClipVisionEmbeddings<B>,
    pre_ln: LayerNorm<B>,
    blocks: Vec<ClipTransformerBlock<B>>,
}

impl<B: Backend> ClipEncoder<B> {
    /// Creates a new CLIP encoder.
    ///
    /// Initializes all layers including embeddings, pre-LayerNorm, and
    /// transformer blocks according to the provided configuration.
    ///
    /// # Arguments
    ///
    /// * `cfg` - CLIP configuration specifying model architecture
    /// * `device` - Device for tensor allocation (CPU/GPU)
    ///
    /// # Returns
    ///
    /// Fully initialized CLIP encoder ready for inference
    pub fn new(cfg: &ClipConfig, device: &B::Device) -> Self {
        info!("Creating CLIP encoder with {} layers", cfg.num_layers);

        let emb = ClipVisionEmbeddings::new(cfg, device);
        let pre_ln = LayerNormConfig::new(cfg.hidden_size)
            .with_epsilon(cfg.layernorm_epsilon as f64)
            .init(device);
        let mut blocks = Vec::new();
        for i in 0..cfg.num_layers {
            debug!(
                "Initializing CLIP transformer block {}/{}",
                i + 1,
                cfg.num_layers
            );
            blocks.push(ClipTransformerBlock::new(cfg, device));
        }

        info!("CLIP encoder created successfully");
        Self {
            emb,
            pre_ln,
            blocks,
        }
    }

    /// Forward pass through CLIP encoder.
    ///
    /// Processes input images through patch embedding, position encoding,
    /// and transformer layers to produce semantic visual features.
    ///
    /// # Arguments
    ///
    /// * `x` - Input image tensor of shape `[B, 3, H, W]`
    /// * `_patch_embeds` - Optional pre-computed patch embeddings (currently unused,
    ///   reserved for future integration with SAM features)
    ///
    /// # Returns
    ///
    /// Feature tensor of shape `[B, num_patches+1, hidden_size]` where the first
    /// token (index 0) is the class token and remaining tokens are patch features.
    ///
    /// # Note
    ///
    /// The `_patch_embeds` parameter is currently ignored. Full integration with
    /// SAM features would require a projection layer to match embedding dimensions.
    pub fn forward(&self, x: Tensor<B, 4>, _patch_embeds: Option<Tensor<B, 4>>) -> Tensor<B, 3> {
        info!("CLIP encoder forward pass: input shape={:?}", x.dims());

        // Note: We ignore patch_embeds for now as dimension matching is complex
        // In full implementation, would need projection layer to match dimensions
        let mut h = self.pre_ln.forward(self.emb.forward(x, None));

        debug!("After embeddings and pre-LN: shape={:?}", h.dims());

        for (i, blk) in self.blocks.iter().enumerate() {
            trace!(
                "Processing CLIP transformer block {}/{}",
                i + 1,
                self.blocks.len()
            );
            h = blk.forward(h);
        }

        info!("CLIP encoder output shape: {:?}", h.dims());
        h
    }

    /// Forward pass excluding class token.
    ///
    /// Convenience method that runs the encoder and strips the class token
    /// from the output sequence, returning only patch feature embeddings.
    ///
    /// # Arguments
    ///
    /// * `x` - Input image tensor of shape `[B, 3, H, W]`
    /// * `patch_embeds` - Optional pre-computed patch embeddings
    ///
    /// # Returns
    ///
    /// Patch features of shape `[B, num_patches, hidden_size]` with the
    /// class token removed (excludes first token in sequence).
    pub fn forward_features(
        &self,
        x: Tensor<B, 4>,
        patch_embeds: Option<Tensor<B, 4>>,
    ) -> Tensor<B, 3> {
        debug!("CLIP encoder forward_features (excluding class token)");
        let h = self.forward(x, patch_embeds);
        let d = h.dims();
        h.slice([0..d[0], 1..d[1], 0..d[2]])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn_ndarray::NdArray;
    type TB = NdArray<f32>;

    #[test]
    fn embedding_shapes() {
        let dev = Default::default();
        let mut cfg = ClipConfig::base();
        cfg.image_size = 32;
        cfg.patch_size = 8;
        cfg.hidden_size = 128;

        let emb = ClipVisionEmbeddings::<TB>::new(&cfg, &dev);
        let out = emb.forward(Tensor::<TB, 4>::zeros([1, 3, 32, 32], &dev), None);

        // 32/8 = 4 patches per side, 4*4 = 16 patches + 1 cls = 17
        assert_eq!(out.dims(), [1, 17, 128]);
    }

    #[test]
    fn shapes_tiny() {
        let dev = Default::default();
        let mut cfg = ClipConfig::base();
        cfg.num_layers = 2;
        cfg.hidden_size = 128;
        cfg.num_attention_heads = 4;
        cfg.ffn_hidden_size = 512;
        cfg.image_size = 32;
        cfg.patch_size = 8;

        let enc = ClipEncoder::<TB>::new(&cfg, &dev);
        let out = enc.forward(Tensor::<TB, 4>::zeros([1, 3, 32, 32], &dev), None);
        assert_eq!(out.dims(), [1, 17, 128]);
    }

    #[test]
    #[ignore]
    fn shapes_full() {
        let dev = Default::default();
        let cfg = ClipConfig::default();
        let enc = ClipEncoder::<TB>::new(&cfg, &dev);
        let out = enc.forward(Tensor::<TB, 4>::zeros([1, 3, 224, 224], &dev), None);
        assert_eq!(out.dims(), [1, 257, 1024]);
    }
}