optical_embeddings/vision/

sam.rs

//! SAM (Segment Anything Model) vision encoder used by Optical Embeddings.
//!
//! This module implements the SAM-based encoder from the DeepSeek-OCR paper, which uses
//! window attention and convolutional compression to efficiently encode images into
//! compact vision tokens.
//!
//! # Architecture
//!
//! The SAM encoder consists of:
//! - **Patch Embedding**: Converts input images to 16×16 patches
//! - **Transformer Blocks**: Process patches with window/global attention
//! - **Neck**: Projects features to output channels (typically 256)
//! - **Compressor**: 16× spatial reduction via two stride-2 conv layers
//!
//! # Example
//!
//! ```ignore
//! let config = SamConfig::base();
//! let encoder = SamEncoder::new(&config, &device);
//! let features = encoder.forward(image_tensor);
//! // Output: [batch, 1024, H/16, W/16] compressed features
//! ```

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

use super::attention::WindowAttention;
use crate::config::SamConfig;

/// Patch embedding layer that converts images to patch tokens.
///
/// Implements a convolutional layer with kernel and stride equal to patch size,
/// effectively partitioning the image into non-overlapping patches.
#[derive(Module, Debug)]
struct PatchEmbed<B: Backend> {
    proj: Conv2d<B>,
    _patch: usize,
    _embed: usize,
}

impl<B: Backend> PatchEmbed<B> {
    /// Creates a new patch embedding layer.
    ///
    /// # Arguments
    /// * `patch_size` - Size of each patch (typically 16)
    /// * `in_chans` - Number of input channels (3 for RGB)
    /// * `embed_dim` - Output embedding dimension
    /// * `device` - Device to create the layer on
    fn new(patch_size: usize, in_chans: usize, embed_dim: usize, device: &B::Device) -> Self {
        debug!(
            "Creating PatchEmbed: patch_size={}, in_chans={}, embed_dim={}",
            patch_size, in_chans, embed_dim
        );
        let proj = Conv2dConfig::new([in_chans, embed_dim], [patch_size, patch_size])
            .with_stride([patch_size, patch_size])
            .with_bias(true)
            .init(device);
        Self {
            proj,
            _patch: patch_size,
            _embed: embed_dim,
        }
    }

    /// Forward pass: converts image to patch embeddings.
    ///
    /// # Shape
    /// - Input: `[B, C, H, W]`
    /// - Output: `[B, embed_dim, H/patch_size, W/patch_size]`
    fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
        trace!("PatchEmbed input shape: {:?}", x.dims());
        let out = self.proj.forward(x);
        trace!("PatchEmbed output shape: {:?}", out.dims());
        out
    }
}

/// Multi-Layer Perceptron (feedforward) block with GELU activation.
///
/// Standard two-layer MLP used in transformer architectures.
#[derive(Module, Debug)]
struct MlpBlock<B: Backend> {
    fc1: Linear<B>,
    fc2: Linear<B>,
}

impl<B: Backend> MlpBlock<B> {
    /// Creates a new MLP block.
    ///
    /// # Arguments
    /// * `dim` - Input/output dimension
    /// * `hidden` - Hidden layer dimension (typically 4× the input dim)
    /// * `device` - Device to create the layer on
    fn new(dim: usize, hidden: usize, device: &B::Device) -> Self {
        debug!("Creating MlpBlock: dim={}, hidden={}", dim, hidden);
        let fc1 = LinearConfig::new(dim, hidden).with_bias(true).init(device);
        let fc2 = LinearConfig::new(hidden, dim).with_bias(true).init(device);
        Self { fc1, fc2 }
    }

    /// Forward pass with GELU activation.
    ///
    /// # Shape
    /// - Input: `[B, N, dim]`
    /// - Output: `[B, N, dim]`
    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        trace!("MlpBlock input shape: {:?}", x.dims());
        let x = self.fc1.forward(x);
        let x = burn::tensor::activation::gelu(x);
        self.fc2.forward(x)
    }
}

/// SAM transformer block with window or global attention.
///
/// Implements a standard transformer block with:
/// - Layer normalization
/// - Window/global multi-head attention
/// - MLP with residual connections
#[derive(Module, Debug)]
struct SamBlock<B: Backend> {
    norm1: LayerNorm<B>,
    attn: WindowAttention<B>,
    norm2: LayerNorm<B>,
    mlp: MlpBlock<B>,
    window_size: usize,
}

impl<B: Backend> SamBlock<B> {
    /// Creates a new SAM transformer block.
    ///
    /// # Arguments
    /// * `dim` - Embedding dimension
    /// * `heads` - Number of attention heads
    /// * `mlp_ratio` - MLP hidden dimension ratio (typically 4.0)
    /// * `window_size` - Window size for attention (0 = global attention)
    /// * `use_rel_pos` - Whether to use relative position bias
    /// * `input_size` - Spatial size of input features
    /// * `device` - Device to create the layer on
    fn new(
        dim: usize,
        heads: usize,
        mlp_ratio: f32,
        window_size: usize,
        use_rel_pos: bool,
        input_size: (usize, usize),
        device: &B::Device,
    ) -> Self {
        debug!(
            "Creating SamBlock: dim={}, heads={}, window_size={}, input_size={:?}",
            dim, heads, window_size, input_size
        );
        let norm1 = LayerNormConfig::new(dim).init(device);
        let norm2 = LayerNormConfig::new(dim).init(device);
        let mlp = MlpBlock::new(dim, (dim as f32 * mlp_ratio) as usize, device);
        let attn = WindowAttention::new(
            dim,
            heads,
            use_rel_pos,
            if window_size == 0 {
                input_size
            } else {
                (window_size, window_size)
            },
            device,
        );
        Self {
            norm1,
            attn,
            norm2,
            mlp,
            window_size,
        }
    }

    /// Forward pass with pre-norm architecture and residual connections.
    ///
    /// # Shape
    /// - Input: `[B, H, W, C]`
    /// - Output: `[B, H, W, C]`
    fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
        trace!("SamBlock input shape: {:?}", x.dims());
        let [b, h, w, c] = {
            let d = x.dims();
            [d[0], d[1], d[2], d[3]]
        };
        let x1 = self
            .norm1
            .forward(x.clone().reshape([b * h * w, c]))
            .reshape([b, h, w, c]);
        let x_attn = if self.window_size > 0 {
            self.attn.forward_windowed(x1, self.window_size)
        } else {
            self.attn.forward(x1)
        };
        let x = x + x_attn;
        let x2 = self
            .norm2
            .forward(x.clone().reshape([b * h * w, c]))
            .reshape([b, h, w, c]);
        let x_mlp = self
            .mlp
            .forward(x2.reshape([b, h * w, c]))
            .reshape([b, h, w, c]);
        x + x_mlp
    }
}

/// SAM vision encoder for the DeepSeek-OCR DeepEncoder architecture.
///
/// Implements the complete SAM encoding pipeline:
/// 1. Patch embedding: Image → patches
/// 2. Position embedding: Add learned position encodings
/// 3. Transformer blocks: Window/global attention processing
/// 4. Neck: Project to output channels
/// 5. Compression: 16× spatial reduction via convolutional layers
///
/// # Compression Pipeline
///
/// For a 1024×1024 image with 16×16 patches:
/// - Input: `[B, 3, 1024, 1024]`
/// - After patch embed: `[B, 768, 64, 64]` (1024/16 = 64)
/// - After transformers: `[B, 768, 64, 64]`
/// - After neck: `[B, 256, 64, 64]`
/// - After comp1: `[B, 512, 32, 32]` (stride=2)
/// - After comp2: `[B, 1024, 16, 16]` (stride=2)
/// - Result: **16× compression** (64² → 16² tokens)
#[derive(Module, Debug)]
pub struct SamEncoder<B: Backend> {
    patch_embed: PatchEmbed<B>,
    pos_embed: Tensor<B, 4>,
    blocks: Vec<SamBlock<B>>,
    neck1: Conv2d<B>,
    neck2: Conv2d<B>,
    comp1: Conv2d<B>,
    comp2: Conv2d<B>,
}

impl<B: Backend> SamEncoder<B> {
    /// Creates a new SAM encoder from configuration.
    ///
    /// # Arguments
    /// * `config` - SAM configuration (see `SamConfig`)
    /// * `device` - Device to create the model on
    ///
    /// # Example
    /// ```ignore
    /// let config = SamConfig::base();
    /// let encoder = SamEncoder::new(&config, &device);
    /// ```
    pub fn new(config: &SamConfig, device: &B::Device) -> Self {
        debug!(
            "Creating PatchEmbed: patch_size={}, in_chans=3, embed_dim={}",
            config.patch_size, config.embed_dim
        );
        let patch_embed = PatchEmbed::new(config.patch_size, 3, config.embed_dim, device);

        let g = config.img_size / config.patch_size;
        info!("Patch grid size: {}x{} = {} patches", g, g, g * g);

        let pos_embed = Tensor::random(
            [1, g, g, config.embed_dim],
            Distribution::Normal(0.0, 0.02),
            device,
        );

        let mut blocks = Vec::new();
        for i in 0..config.depth {
            if config.global_attn_indexes.contains(&i) {
                debug!("Layer {}: using global attention", i);
            } else {
                debug!(
                    "Layer {}: using window attention (size={})",
                    i, config.window_size
                );
            }

            let window = if config.global_attn_indexes.contains(&i) {
                0
            } else {
                config.window_size
            };
            blocks.push(SamBlock::new(
                config.embed_dim,
                config.num_heads,
                config.mlp_ratio,
                window,
                config.use_rel_pos,
                (g, g),
                device,
            ));
        }

        debug!(
            "Creating neck layers: embed_dim={} -> out_chans={}",
            config.embed_dim, config.out_chans
        );
        let neck1 = Conv2dConfig::new([config.embed_dim, config.out_chans], [1, 1])
            .with_bias(false)
            .init(device);
        let neck2 = Conv2dConfig::new([config.out_chans, config.out_chans], [3, 3])
            .with_padding(PaddingConfig2d::Same)
            .with_bias(false)
            .init(device);

        debug!(
            "Creating compression layers: {} -> 512 -> 1024 (stride=2 each)",
            config.out_chans
        );

        // Use Explicit padding to allow spatial reduction with stride=2
        let comp1 = Conv2dConfig::new([config.out_chans, 512], [3, 3])
            .with_stride([2, 2])
            .with_padding(PaddingConfig2d::Explicit(1, 1))
            .with_bias(false)
            .init(device);

        let comp2 = Conv2dConfig::new([512, 1024], [3, 3])
            .with_stride([2, 2])
            .with_padding(PaddingConfig2d::Explicit(1, 1))
            .with_bias(false)
            .init(device);

        Self {
            patch_embed,
            pos_embed,
            blocks,
            neck1,
            neck2,
            comp1,
            comp2,
        }
    }

    /// Encodes an image into compressed vision features.
    ///
    /// # Arguments
    /// * `x` - Input image tensor
    ///
    /// # Shape
    /// - Input: `[B, 3, H, W]`
    /// - Output: `[B, 1024, H/P/4, W/P/4]` where P is patch_size
    ///
    /// For default settings (patch_size=16, 16× compression):
    /// - 1024×1024 image → `[B, 1024, 16, 16]` (256 tokens)
    /// - 512×512 image → `[B, 1024, 8, 8]` (64 tokens)
    ///
    /// # Example
    /// ```ignore
    /// let image = Tensor::zeros(, &device);[1]
    /// let features = encoder.forward(image);
    /// // Shape:  = 256 vision tokens[2][1]
    /// ```
    pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
        info!("SamEncoder forward: input shape {:?}", x.dims());

        let x = self.patch_embed.forward(x); // [B, C, H', W']
        debug!("After patch_embed: {:?}", x.dims());

        let x = x.swap_dims(1, 3).swap_dims(1, 2); // [B, H', W', C]
        debug!("After transpose: {:?}", x.dims());

        let x = x.clone() + self.interpolate_pos_embed(&x);
        debug!("After pos_embed: {:?}", x.dims());

        let mut x = x;
        for (i, blk) in self.blocks.iter().enumerate() {
            x = blk.forward(x);
            trace!("After block {}: {:?}", i, x.dims());
        }
        info!(
            "After all {} transformer blocks: {:?}",
            self.blocks.len(),
            x.dims()
        );

        let x = x.swap_dims(1, 3).swap_dims(2, 3); // [B, C, H', W']
        debug!("After transpose back: {:?}", x.dims());

        let x = self.neck1.forward(x);
        debug!("After neck1: {:?}", x.dims());

        let x = self.neck2.forward(x);
        debug!("After neck2: {:?}", x.dims());

        let x = self.comp1.forward(x);
        info!("After comp1 (stride=2): {:?}", x.dims());

        let out = self.comp2.forward(x);
        info!("After comp2 (stride=2): {:?}", out.dims());
        info!(
            "Final compression: {}x{} = {} tokens",
            out.dims()[2],
            out.dims()[3],
            out.dims()[2] * out.dims()[3]
        );

        out // [B, 1024, H'', W'']
    }

    /// Interpolates position embeddings to match input size.
    ///
    /// Currently returns zeros; full implementation would use bicubic interpolation
    /// for dynamic resolution support.
    fn interpolate_pos_embed(&self, x: &Tensor<B, 4>) -> Tensor<B, 4> {
        let d = x.dims();
        Tensor::zeros([d[0], d[1], d[2], d[3]], &x.device())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn_ndarray::NdArray;
    type TB = NdArray<f32>;
    #[test]
    fn shape_sanity() {
        let dev = Default::default();
        let mut cfg = SamConfig::default();
        cfg.img_size = 1024;
        let enc = SamEncoder::<TB>::new(&cfg, &dev);
        let out = enc.forward(Tensor::<TB, 4>::zeros([1, 3, 1024, 1024], &dev));
        assert_eq!(out.dims()[1], 1024);
    }
}