Upload folder using huggingface_hub

configuration_latex_decoder.py

@@ -1,48 +1,48 @@
-from transformers import PretrainedConfig
-
-
-class LaTeXDecoderConfig(PretrainedConfig):
-    model_type = "latex_decoder"
-
-    def __init__(
-        self,
-        vocab_size: int = 8192,
-        pad_id: int = 0,
-        bos_id: int = 2,
-        eos_id: int = 3,
-        d_model: int = 512,
-        n_heads: int = 8,
-        n_layers: int = 6,
-        d_ff: int = 1408,
-        dropout: float = 0.1,
-        max_seq_len: int = 200,
-        rope_theta: float = 10000.0,
-        tie_weights: bool = False,
-        **kwargs,
-    ):
-        kwargs.pop("pad_token_id", None)
-        kwargs.pop("bos_token_id", None)
-        kwargs.pop("eos_token_id", None)
-        super().__init__(
-            pad_token_id=pad_id,
-            bos_token_id=bos_id,
-            eos_token_id=eos_id,
-            **kwargs,
-        )
-        self.vocab_size = vocab_size
-        self.pad_id = pad_id
-        self.bos_id = bos_id
-        self.eos_id = eos_id
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.n_layers = n_layers
-        self.d_ff = d_ff
-        self.dropout = dropout
-        self.max_seq_len = max_seq_len
-        self.rope_theta = rope_theta
-        self.tie_weights = tie_weights
-
-    @property
-    def head_dim(self) -> int:
-        assert self.d_model % self.n_heads == 0
-        return self.d_model // self.n_heads
+from transformers import PretrainedConfig
+
+
+class LaTeXDecoderConfig(PretrainedConfig):
+    model_type = "latex_decoder"
+
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        pad_id: int = 0,
+        bos_id: int = 2,
+        eos_id: int = 3,
+        d_model: int = 512,
+        n_heads: int = 8,
+        n_layers: int = 6,
+        d_ff: int = 1408,
+        dropout: float = 0.1,
+        max_seq_len: int = 200,
+        rope_theta: float = 10000.0,
+        tie_weights: bool = False,
+        **kwargs,
+    ):
+        kwargs.pop("pad_token_id", None)
+        kwargs.pop("bos_token_id", None)
+        kwargs.pop("eos_token_id", None)
+        super().__init__(
+            pad_token_id=pad_id,
+            bos_token_id=bos_id,
+            eos_token_id=eos_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.pad_id = pad_id
+        self.bos_id = bos_id
+        self.eos_id = eos_id
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.d_ff = d_ff
+        self.dropout = dropout
+        self.max_seq_len = max_seq_len
+        self.rope_theta = rope_theta
+        self.tie_weights = tie_weights
+
+    @property
+    def head_dim(self) -> int:
+        assert self.d_model % self.n_heads == 0
+        return self.d_model // self.n_heads

configuration_latex_ocr.py

@@ -1,66 +1,66 @@
-from transformers import PretrainedConfig
-
-
-class Nav2TexConfig(PretrainedConfig):
-    model_type = "nav2tex"
-
-    def __init__(
-        self,
-        patch_size: int = 16,
-        image_height: int = 64,
-        max_image_width: int = 1024,
-        max_image_height: int = 640,
-        resize_in_dataset: bool = True,
-        max_token_len: int = 200,
-        navit_dim: int = 512,
-        navit_depth: int = 8,
-        navit_heads: int = 8,
-        navit_dim_head: int = 64,
-        navit_mlp_dim: int = 2048,
-        navit_dropout: float = 0.0,
-        navit_emb_dropout: float = 0.0,
-        vision_hidden_size: int = 512,
-        llm_hidden_size: int = 512,
-        projector_intermediate_size: int = 1024,
-        max_visual_tokens: int = 256,
-        max_new_tokens: int = 200,
-        num_beams: int = 4,
-        decoder_arch: dict | None = None,
-        decoder_weights_tied: bool = False,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.patch_size = patch_size
-        self.image_height = image_height
-        self.max_image_width = max_image_width
-        self.max_image_height = max_image_height
-        self.resize_in_dataset = resize_in_dataset
-        self.max_token_len = max_token_len
-        self.navit_dim = navit_dim
-        self.navit_depth = navit_depth
-        self.navit_heads = navit_heads
-        self.navit_dim_head = navit_dim_head
-        self.navit_mlp_dim = navit_mlp_dim
-        self.navit_dropout = navit_dropout
-        self.navit_emb_dropout = navit_emb_dropout
-        self.vision_hidden_size = vision_hidden_size
-        self.llm_hidden_size = llm_hidden_size
-        self.projector_intermediate_size = projector_intermediate_size
-        self.max_visual_tokens = max_visual_tokens
-        self.max_new_tokens = max_new_tokens
-        self.num_beams = num_beams
-        self.decoder_arch = decoder_arch or {
-            "vocab_size": 2046,
-            "pad_id": 0,
-            "bos_id": 2,
-            "eos_id": 3,
-            "d_model": 512,
-            "n_heads": 8,
-            "n_layers": 6,
-            "d_ff": 1408,
-            "dropout": 0.1,
-            "max_seq_len": 200,
-            "rope_theta": 10000.0,
-            "tie_weights": True,
-        }
-        self.decoder_weights_tied = decoder_weights_tied
+from transformers import PretrainedConfig
+
+
+class Nav2TexConfig(PretrainedConfig):
+    model_type = "nav2tex"
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        image_height: int = 64,
+        max_image_width: int = 1024,
+        max_image_height: int = 640,
+        resize_in_dataset: bool = True,
+        max_token_len: int = 200,
+        navit_dim: int = 512,
+        navit_depth: int = 8,
+        navit_heads: int = 8,
+        navit_dim_head: int = 64,
+        navit_mlp_dim: int = 2048,
+        navit_dropout: float = 0.0,
+        navit_emb_dropout: float = 0.0,
+        vision_hidden_size: int = 512,
+        llm_hidden_size: int = 512,
+        projector_intermediate_size: int = 1024,
+        max_visual_tokens: int = 256,
+        max_new_tokens: int = 200,
+        num_beams: int = 4,
+        decoder_arch: dict | None = None,
+        decoder_weights_tied: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.patch_size = patch_size
+        self.image_height = image_height
+        self.max_image_width = max_image_width
+        self.max_image_height = max_image_height
+        self.resize_in_dataset = resize_in_dataset
+        self.max_token_len = max_token_len
+        self.navit_dim = navit_dim
+        self.navit_depth = navit_depth
+        self.navit_heads = navit_heads
+        self.navit_dim_head = navit_dim_head
+        self.navit_mlp_dim = navit_mlp_dim
+        self.navit_dropout = navit_dropout
+        self.navit_emb_dropout = navit_emb_dropout
+        self.vision_hidden_size = vision_hidden_size
+        self.llm_hidden_size = llm_hidden_size
+        self.projector_intermediate_size = projector_intermediate_size
+        self.max_visual_tokens = max_visual_tokens
+        self.max_new_tokens = max_new_tokens
+        self.num_beams = num_beams
+        self.decoder_arch = decoder_arch or {
+            "vocab_size": 2046,
+            "pad_id": 0,
+            "bos_id": 2,
+            "eos_id": 3,
+            "d_model": 512,
+            "n_heads": 8,
+            "n_layers": 6,
+            "d_ff": 1408,
+            "dropout": 0.1,
+            "max_seq_len": 200,
+            "rope_theta": 10000.0,
+            "tie_weights": True,
+        }
+        self.decoder_weights_tied = decoder_weights_tied

modeling_latex_decoder.py

@@ -1,202 +1,202 @@
-# update v2
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional
-
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import CausalLMOutput
-
-try:
-    from .configuration_latex_decoder import LaTeXDecoderConfig
-except ImportError:
-    from latex_ocr.configuration_latex_decoder import LaTeXDecoderConfig
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, d_model: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(d_model))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
-        return x / rms * self.weight
-
-
-def _build_rope_cache(seq_len, head_dim, theta=10000.0, device=None, dtype=torch.float32):
-    half = head_dim // 2
-    inv_freq = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
-    pos = torch.arange(seq_len, device=device, dtype=torch.float32)
-    freqs = torch.outer(pos, inv_freq)
-    emb = torch.cat([freqs, freqs], dim=-1)
-    return emb.cos().to(dtype), emb.sin().to(dtype)
-
-
-def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-    half = x.shape[-1] // 2
-    x1, x2 = x[..., :half], x[..., half:]
-    return torch.cat([-x2, x1], dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos.unsqueeze(0).unsqueeze(0)
-    sin = sin.unsqueeze(0).unsqueeze(0)
-    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin
-
-
-class CausalSelfAttention(nn.Module):
-    def __init__(self, cfg: LaTeXDecoderConfig):
-        super().__init__()
-        self.n_heads = cfg.n_heads
-        self.head_dim = cfg.head_dim
-        self.d_model = cfg.d_model
-        self.dropout_p = cfg.dropout
-        self.rope_theta = cfg.rope_theta
-
-        self.qkv_proj = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
-        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
-        self._rope_cache: dict = {}
-
-    def _get_rope(self, seq_len, device, dtype):
-        key = (seq_len, str(device), dtype)
-        if key not in self._rope_cache:
-            self._rope_cache[key] = _build_rope_cache(seq_len, self.head_dim, self.rope_theta, device, dtype)
-        return self._rope_cache[key]
-
-    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        B, T, C = x.shape
-        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
-
-        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-
-        cos, sin = self._get_rope(T, x.device, q.dtype)
-        q, k = apply_rope(q, k, cos, sin)
-
-        dropout_p = self.dropout_p if self.training else 0.0
-
-        if attention_mask is not None:
-            causal = torch.triu(torch.full((T, T), float("-inf"), device=x.device, dtype=q.dtype), diagonal=1)
-            pad = (~attention_mask).unsqueeze(1).unsqueeze(2)
-            attn_bias = causal.unsqueeze(0).unsqueeze(0).expand(B, 1, T, T).clone()
-            attn_bias = attn_bias.masked_fill(pad, float("-inf"))
-            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias, dropout_p=dropout_p, is_causal=False)
-        else:
-            out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p, is_causal=True)
-
-        return self.out_proj(out.transpose(1, 2).contiguous().view(B, T, C))
-
-
-class SwiGLUFFN(nn.Module):
-    def __init__(self, cfg: LaTeXDecoderConfig):
-        super().__init__()
-        self.gate_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
-        self.up_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
-        self.down_proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
-        self.dropout = nn.Dropout(cfg.dropout)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, cfg: LaTeXDecoderConfig):
-        super().__init__()
-        self.norm1 = RMSNorm(cfg.d_model)
-        self.attn = CausalSelfAttention(cfg)
-        self.norm2 = RMSNorm(cfg.d_model)
-        self.ffn = SwiGLUFFN(cfg)
-        self.drop = nn.Dropout(cfg.dropout)
-
-    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        x = x + self.drop(self.attn(self.norm1(x), attention_mask))
-        x = x + self.drop(self.ffn(self.norm2(x)))
-        return x
-
-
-class LaTeXDecoderForCausalLM(PreTrainedModel):
-    config_class = LaTeXDecoderConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = False
-
-    def __init__(self, config: LaTeXDecoderConfig):
-        super().__init__(config)
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, padding_idx=config.pad_id)
-        self.embed_drop = nn.Dropout(config.dropout)
-        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
-        self.norm_final = RMSNorm(config.d_model)
-        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
-
-        if config.tie_weights:
-            self.lm_head.weight = self.embed_tokens.weight
-
-        self.post_init()
-
-    def _init_weights(self, module: nn.Module):
-        if isinstance(module, nn.Linear):
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> CausalLMOutput:
-        x = self.embed_drop(self.embed_tokens(input_ids))
-        for layer in self.layers:
-            x = layer(x, attention_mask)
-        logits = self.lm_head(self.norm_final(x))
-
-        loss = None
-        if labels is not None:
-            shift_logits = logits[:, :-1, :].contiguous()
-            shift_labels = labels[:, 1:].contiguous()
-            shift_labels = shift_labels.masked_fill(shift_labels == self.config.pad_id, -100)
-            loss = F.cross_entropy(
-                shift_logits.view(-1, self.config.vocab_size),
-                shift_labels.view(-1),
-                ignore_index=-100,
-            )
-
-        return CausalLMOutput(loss=loss, logits=logits)
-
-    @torch.inference_mode()
-    def generate(
-        self,
-        prompt_ids: torch.Tensor,
-        max_new_tokens: int = 200,
-        temperature: float = 1.0,
-        top_p: float = 0.9,
-        eos_id: Optional[int] = None,
-    ) -> torch.Tensor:
-        eos = eos_id if eos_id is not None else self.config.eos_id
-        generated = prompt_ids.clone()
-
-        for _ in range(max_new_tokens):
-            ctx = generated[:, -self.config.max_seq_len:]
-            logits = self.forward(ctx).logits[:, -1, :]
-
-            if temperature == 0.0:
-                next_id = logits.argmax(dim=-1, keepdim=True)
-            else:
-                probs = F.softmax(logits / temperature, dim=-1)
-                sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
-                cumsum = sorted_probs.cumsum(dim=-1)
-                sorted_probs[cumsum - sorted_probs > top_p] = 0.0
-                sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)
-                next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
-
-            generated = torch.cat([generated, next_id], dim=-1)
-            if next_id.item() == eos:
-                break
-
-        return generated
+# update v2
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
+try:
+    from .configuration_latex_decoder import LaTeXDecoderConfig
+except ImportError:
+    from latex_ocr.configuration_latex_decoder import LaTeXDecoderConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return x / rms * self.weight
+
+
+def _build_rope_cache(seq_len, head_dim, theta=10000.0, device=None, dtype=torch.float32):
+    half = head_dim // 2
+    inv_freq = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
+    pos = torch.arange(seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(pos, inv_freq)
+    emb = torch.cat([freqs, freqs], dim=-1)
+    return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(q, k, cos, sin):
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: LaTeXDecoderConfig):
+        super().__init__()
+        self.n_heads = cfg.n_heads
+        self.head_dim = cfg.head_dim
+        self.d_model = cfg.d_model
+        self.dropout_p = cfg.dropout
+        self.rope_theta = cfg.rope_theta
+
+        self.qkv_proj = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self._rope_cache: dict = {}
+
+    def _get_rope(self, seq_len, device, dtype):
+        key = (seq_len, str(device), dtype)
+        if key not in self._rope_cache:
+            self._rope_cache[key] = _build_rope_cache(seq_len, self.head_dim, self.rope_theta, device, dtype)
+        return self._rope_cache[key]
+
+    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, T, C = x.shape
+        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
+
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self._get_rope(T, x.device, q.dtype)
+        q, k = apply_rope(q, k, cos, sin)
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        if attention_mask is not None:
+            causal = torch.triu(torch.full((T, T), float("-inf"), device=x.device, dtype=q.dtype), diagonal=1)
+            pad = (~attention_mask).unsqueeze(1).unsqueeze(2)
+            attn_bias = causal.unsqueeze(0).unsqueeze(0).expand(B, 1, T, T).clone()
+            attn_bias = attn_bias.masked_fill(pad, float("-inf"))
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias, dropout_p=dropout_p, is_causal=False)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p, is_causal=True)
+
+        return self.out_proj(out.transpose(1, 2).contiguous().view(B, T, C))
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, cfg: LaTeXDecoderConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.up_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.down_proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
+        self.dropout = nn.Dropout(cfg.dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg: LaTeXDecoderConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(cfg.d_model)
+        self.attn = CausalSelfAttention(cfg)
+        self.norm2 = RMSNorm(cfg.d_model)
+        self.ffn = SwiGLUFFN(cfg)
+        self.drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x = x + self.drop(self.attn(self.norm1(x), attention_mask))
+        x = x + self.drop(self.ffn(self.norm2(x)))
+        return x
+
+
+class LaTeXDecoderForCausalLM(PreTrainedModel):
+    config_class = LaTeXDecoderConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: LaTeXDecoderConfig):
+        super().__init__(config)
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, padding_idx=config.pad_id)
+        self.embed_drop = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.norm_final = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        if config.tie_weights:
+            self.lm_head.weight = self.embed_tokens.weight
+
+        self.post_init()
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> CausalLMOutput:
+        x = self.embed_drop(self.embed_tokens(input_ids))
+        for layer in self.layers:
+            x = layer(x, attention_mask)
+        logits = self.lm_head(self.norm_final(x))
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            shift_labels = shift_labels.masked_fill(shift_labels == self.config.pad_id, -100)
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 200,
+        temperature: float = 1.0,
+        top_p: float = 0.9,
+        eos_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        eos = eos_id if eos_id is not None else self.config.eos_id
+        generated = prompt_ids.clone()
+
+        for _ in range(max_new_tokens):
+            ctx = generated[:, -self.config.max_seq_len:]
+            logits = self.forward(ctx).logits[:, -1, :]
+
+            if temperature == 0.0:
+                next_id = logits.argmax(dim=-1, keepdim=True)
+            else:
+                probs = F.softmax(logits / temperature, dim=-1)
+                sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
+                cumsum = sorted_probs.cumsum(dim=-1)
+                sorted_probs[cumsum - sorted_probs > top_p] = 0.0
+                sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)
+                next_id = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
+
+            generated = torch.cat([generated, next_id], dim=-1)
+            if next_id.item() == eos:
+                break
+
+        return generated

modeling_latex_ocr.py

@@ -1,508 +1,508 @@
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from functools import partial
-from torch import nn
-from torch.nn.utils.rnn import pad_sequence as orig_pad_sequence
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput
-
-try:
-    from .configuration_latex_decoder import LaTeXDecoderConfig
-    from .configuration_latex_ocr import Nav2TexConfig
-    from .modeling_latex_decoder import LaTeXDecoderForCausalLM
-except ImportError:
-    from nav2tex.configuration_latex_decoder import LaTeXDecoderConfig
-    from nav2tex.configuration_latex_ocr import Nav2TexConfig
-    from nav2tex.modeling_latex_decoder import LaTeXDecoderForCausalLM
-
-try:
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import pad_input, unpad_input
-    HAS_FLASH_ATTN = True
-except ImportError:
-    HAS_FLASH_ATTN = False
-
-
-def exists(val):
-    return val is not None
-
-
-def divisible_by(numer, denom):
-    return (numer % denom) == 0
-
-
-class LayerNorm(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.normalized_shape = (dim,)
-        self.eps = 1e-5
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.bias = nn.Parameter(torch.zeros(dim))
-
-    def forward(self, x):
-        return F.layer_norm(
-            x.float(), self.normalized_shape,
-            self.weight.float(), self.bias.float(), self.eps,
-        ).to(x.dtype)
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, heads, dim):
-        super().__init__()
-        self.scale = dim ** 0.5
-        self.gamma = nn.Parameter(torch.ones(heads, 1, dim))
-
-    def forward(self, x):
-        return F.normalize(x, dim=-1) * self.scale * self.gamma.to(x.dtype)
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat([-x2, x1], dim=-1)
-
-
-def apply_2d_rope(q, k, h_idx, w_idx):
-    _, _, _, d = q.shape
-    if d % 4 != 0:
-        raise ValueError(f"apply_2d_rope expects dim_head divisible by 4, got D={d}")
-    dim_half = d // 2
-    dim_quarter = d // 4
-    inv_freq = 1.0 / (10000 ** (torch.arange(dim_quarter, device=q.device).float() / dim_quarter))
-    h_theta = h_idx[..., None].float() * inv_freq
-    w_theta = w_idx[..., None].float() * inv_freq
-    sin_h = torch.cat([h_theta.sin(), h_theta.sin()], dim=-1).to(q.dtype)[:, None, :, :]
-    cos_h = torch.cat([h_theta.cos(), h_theta.cos()], dim=-1).to(q.dtype)[:, None, :, :]
-    sin_w = torch.cat([w_theta.sin(), w_theta.sin()], dim=-1).to(q.dtype)[:, None, :, :]
-    cos_w = torch.cat([w_theta.cos(), w_theta.cos()], dim=-1).to(q.dtype)[:, None, :, :]
-
-    def rope(x, sin, cos):
-        return x * cos + rotate_half(x) * sin
-
-    q = torch.cat([rope(q[..., :dim_half], sin_h, cos_h), rope(q[..., dim_half:], sin_w, cos_w)], dim=-1)
-    k = torch.cat([rope(k[..., :dim_half], sin_h, cos_h), rope(k[..., dim_half:], sin_w, cos_w)], dim=-1)
-    return q, k
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, hidden_dim, dropout=0.0):
-        super().__init__()
-        self.net = nn.Sequential(
-            LayerNorm(dim),
-            nn.Linear(dim, hidden_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, dim),
-            nn.Dropout(dropout),
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
-        super().__init__()
-        inner_dim = dim_head * heads
-        self.heads = heads
-        self.norm = LayerNorm(dim)
-        self.q_norm = RMSNorm(heads, dim_head)
-        self.k_norm = RMSNorm(heads, dim_head)
-        self.to_q = nn.Linear(dim, inner_dim, bias=False)
-        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
-        self.attend = nn.Softmax(dim=-1)
-        self.dropout = nn.Dropout(dropout)
-        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim, bias=False), nn.Dropout(dropout))
-
-    def forward(self, x, mask=None, attn_mask=None, positions=None):
-        x = self.norm(x)
-        q = self.to_q(x)
-        k, v = self.to_kv(x).chunk(2, dim=-1)
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (q, k, v))
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        if positions is not None:
-            q, k = apply_2d_rope(q, k, positions[0], positions[1])
-
-        if HAS_FLASH_ATTN and x.is_cuda and attn_mask is None:
-            fa_dtype = q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
-            q_ = rearrange(q, "b h n d -> b n h d").contiguous().to(fa_dtype)
-            k_ = rearrange(k, "b h n d -> b n h d").contiguous().to(fa_dtype)
-            v_ = rearrange(v, "b h n d -> b n h d").contiguous().to(fa_dtype)
-            if exists(mask):
-                batch, seqlen = mask.shape
-                q_unpad, indices, cu_q, max_q, *_ = unpad_input(q_, mask)
-                k_unpad, _, cu_k, max_k, *_ = unpad_input(k_, mask)
-                v_unpad, _, _, _, *_ = unpad_input(v_, mask)
-                out_unpad = flash_attn_varlen_func(
-                    q_unpad, k_unpad, v_unpad,
-                    cu_seqlens_q=cu_q, cu_seqlens_k=cu_k,
-                    max_seqlen_q=max_q, max_seqlen_k=max_k,
-                    dropout_p=self.dropout.p if self.training else 0.0,
-                    causal=False,
-                )
-                out = pad_input(out_unpad, indices, batch, seqlen)
-            else:
-                out = flash_attn_func(
-                    q_, k_, v_,
-                    dropout_p=self.dropout.p if self.training else 0.0,
-                    causal=False,
-                )
-            out = rearrange(out, "b n h d -> b n (h d)").to(x.dtype)
-        else:
-            dots = torch.matmul(q, k.transpose(-1, -2))
-            if exists(mask):
-                dots = dots.masked_fill(~mask[:, None, None, :], -torch.finfo(dots.dtype).max)
-            if exists(attn_mask):
-                dots = dots.masked_fill(~attn_mask, -torch.finfo(dots.dtype).max)
-            attn = self.dropout(self.attend(dots))
-            out = rearrange(torch.matmul(attn, v), "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class Transformer(nn.Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.0):
-        super().__init__()
-        self.layers = nn.ModuleList([
-            nn.ModuleList([Attention(dim, heads, dim_head, dropout), FeedForward(dim, mlp_dim, dropout)])
-            for _ in range(depth)
-        ])
-        self.norm = LayerNorm(dim)
-
-    def forward(self, x, mask=None, attn_mask=None, positions=None):
-        for attn, ff in self.layers:
-            x = attn(x, mask=mask, attn_mask=attn_mask, positions=positions) + x
-            x = ff(x) + x
-        return self.norm(x)
-
-
-class NaViT_Encoder(nn.Module):
-    def __init__(self, *, image_size, patch_size, dim, depth, heads, mlp_dim,
-                 channels=3, dim_head=64, dropout=0.0, emb_dropout=0.0):
-        super().__init__()
-        image_height, image_width = image_size
-        assert divisible_by(image_height, patch_size)
-        assert divisible_by(image_width, patch_size)
-        self.patch_size = patch_size
-        self.to_patch_embedding = nn.Sequential(
-            LayerNorm(channels * patch_size ** 2),
-            nn.Linear(channels * patch_size ** 2, dim),
-            LayerNorm(dim),
-        )
-        self.dropout = nn.Dropout(emb_dropout)
-        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    def forward(self, batched_images):
-        p = self.patch_size
-        device = self.device
-        arange = partial(torch.arange, device=device)
-        pad_sequence = partial(orig_pad_sequence, batch_first=True)
-        batched_sequences, batched_positions = [], []
-
-        for images in batched_images:
-            sequences, positions = [], []
-            for image in images:
-                _, h, w = image.shape
-                ph, pw = h // p, w // p
-                seq = rearrange(image, "c (h p1) (w p2) -> (h w) (c p1 p2)", p1=p, p2=p)
-                pos = torch.stack(torch.meshgrid(arange(ph), arange(pw), indexing="ij"), dim=-1)
-                sequences.append(seq)
-                positions.append(rearrange(pos, "h w c -> (h w) c"))
-            batched_sequences.append(torch.cat(sequences, dim=0))
-            batched_positions.append(torch.cat(positions, dim=0))
-
-        patches = pad_sequence(batched_sequences)
-        patch_positions = pad_sequence(batched_positions)
-        lengths = torch.tensor([seq.shape[0] for seq in batched_sequences], device=device)
-        mask = torch.arange(patches.shape[1], device=device)[None, :] < lengths[:, None]
-        x = self.to_patch_embedding(patches.to(next(self.parameters()).dtype))
-        h_idx, w_idx = patch_positions.unbind(dim=-1)
-        x = self.dropout(x)
-        x = self.transformer(x, mask=mask, positions=(h_idx, w_idx))
-        return x, mask
-
-
-class MLPProjector(nn.Module):
-    def __init__(self, vision_hidden_size=1024, llm_hidden_size=512, intermediate_size=2048):
-        super().__init__()
-        self.norm = nn.LayerNorm(vision_hidden_size)
-        self.gate_proj = nn.Linear(vision_hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(vision_hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, llm_hidden_size, bias=False)
-
-    def forward(self, x):
-        x = self.norm(x)
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class VisualEncoder(nn.Module):
-    def __init__(self, encoder, bridge, max_visual_tokens):
-        super().__init__()
-        self.navit = encoder
-        self.projector = bridge
-        self.max_visual_tokens = max_visual_tokens
-
-    def forward(self, batched_images):
-        x, mask = self.navit(batched_images)
-        if x.shape[1] > self.max_visual_tokens:
-            x = x[:, :self.max_visual_tokens]
-            mask = mask[:, :self.max_visual_tokens]
-        return self.projector(x), mask
-
-
-class CustomDecoder(nn.Module):
-    def __init__(self, config: Nav2TexConfig):
-        super().__init__()
-        dec = config.decoder_arch
-        self._model = LaTeXDecoderForCausalLM(
-            LaTeXDecoderConfig(
-                vocab_size=dec["vocab_size"],
-                pad_id=dec["pad_id"],
-                bos_id=dec["bos_id"],
-                eos_id=dec["eos_id"],
-                d_model=dec["d_model"],
-                n_heads=dec["n_heads"],
-                n_layers=dec["n_layers"],
-                d_ff=dec["d_ff"],
-                dropout=dec.get("dropout", 0.1),
-                max_seq_len=dec["max_seq_len"],
-                rope_theta=dec.get("rope_theta", 10000.0),
-                tie_weights=dec.get("tie_weights", True),
-            )
-        )
-        self.pad_token_id = self._model.config.pad_id
-        self.eos_token_id = self._model.config.eos_id
-        self._vocab_size = self._model.config.vocab_size
-        self._pad_id = self._model.config.pad_id
-        if not config.decoder_weights_tied:
-            self.untie_weights()
-
-    def get_input_embeddings(self):
-        return self._model.embed_tokens
-
-    def tie_weights(self):
-        self._model.lm_head.weight = self._model.embed_tokens.weight
-
-    def untie_weights(self):
-        if self.are_weights_tied():
-            self._model.lm_head.weight = nn.Parameter(self._model.embed_tokens.weight.detach().clone())
-
-    def are_weights_tied(self):
-        return self._model.lm_head.weight.data_ptr() == self._model.embed_tokens.weight.data_ptr()
-
-    def _forward_embeds(self, inputs_embeds, attention_mask=None):
-        x = self._model.embed_drop(inputs_embeds)
-        mask = attention_mask.bool() if attention_mask is not None else None
-        for layer in self._model.layers:
-            x = layer(x, mask)
-        return self._model.lm_head(self._model.norm_final(x))
-
-    def forward(self, inputs_embeds=None, attention_mask=None, labels=None, **kwargs):
-        logits = self._forward_embeds(inputs_embeds, attention_mask)
-        loss = None
-        if labels is not None:
-            shift_logits = logits[:, :-1].contiguous()
-            shift_labels = labels[:, 1:].contiguous().masked_fill(
-                labels[:, 1:].contiguous() == self._pad_id, -100
-            )
-            loss = F.cross_entropy(
-                shift_logits.view(-1, self._vocab_size),
-                shift_labels.view(-1),
-                ignore_index=-100,
-            )
-        return BaseModelOutput(last_hidden_state=logits, hidden_states=(loss,))
-
-    @torch.no_grad()
-    def generate(self, inputs_embeds, attention_mask, max_new_tokens, num_beams=1):
-        device = inputs_embeds.device
-        batch  = inputs_embeds.shape[0]
-
-        if num_beams > 1:
-            # beam search: only supports batch_size=1
-            assert batch == 1, "beam search only supports batch_size=1"
-            return self._beam_search(inputs_embeds, attention_mask, max_new_tokens, num_beams)
-
-        return self._greedy_batch(inputs_embeds, attention_mask, max_new_tokens)
-
-    @torch.no_grad()
-    def _greedy_batch(self, inputs_embeds, attention_mask, max_new_tokens):
-        """Greedy decoding with true batch support."""
-        eos_id  = self.eos_token_id
-        pad_id  = self._pad_id
-        device  = inputs_embeds.device
-        batch   = inputs_embeds.shape[0]
-        d_model = inputs_embeds.shape[-1]
-
-        # generated token ids per sample, and finished flags
-        gen_ids  = [[] for _ in range(batch)]
-        finished = torch.zeros(batch, dtype=torch.bool, device=device)
-
-        cur_embeds = inputs_embeds                          # (B, vis_len, D)
-        cur_mask   = attention_mask                         # (B, vis_len)
-
-        for _ in range(max_new_tokens):
-            logits   = self._forward_embeds(cur_embeds, cur_mask)  # (B, seq, vocab)
-            next_tok = logits[:, -1, :].argmax(dim=-1)             # (B,)
-
-            for i in range(batch):
-                if not finished[i]:
-                    gen_ids[i].append(next_tok[i].item())
-            finished |= (next_tok == eos_id)
-            if finished.all():
-                break
-
-            tok_emb  = self._model.embed_tokens(next_tok.unsqueeze(1))   # (B, 1, D)
-            tok_mask = cur_mask.new_ones(batch, 1)
-            cur_embeds = torch.cat([cur_embeds, tok_emb], dim=1)
-            cur_mask   = torch.cat([cur_mask,   tok_mask], dim=1)
-
-        # pad to same length and return (B, max_len)
-        max_len = max((len(ids) for ids in gen_ids), default=0)
-        if max_len == 0:
-            return torch.zeros(batch, 0, dtype=torch.long, device=device)
-        out = torch.full((batch, max_len), pad_id, dtype=torch.long, device=device)
-        for i, ids in enumerate(gen_ids):
-            if ids:
-                out[i, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device)
-        return out
-
-    @torch.no_grad()
-    def _beam_search(self, inputs_embeds, attention_mask, max_new_tokens, num_beams):
-        """Original beam search (batch_size=1 only)."""
-        eos_id   = self.eos_token_id
-        device   = inputs_embeds.device
-        vis_emb  = inputs_embeds[0]
-        vis_len  = vis_emb.shape[0]
-        vis_mask = attention_mask[0] if attention_mask is not None else None
-        beams    = [(0.0, [], False) for _ in range(num_beams)]
-
-        for _ in range(max_new_tokens):
-            all_embeds, all_masks = [], []
-            for score, ids, _ in beams:
-                tok_emb = self._model.embed_tokens(torch.tensor(ids, device=device, dtype=torch.long)) if ids else None
-                seq_emb = torch.cat([vis_emb, tok_emb], dim=0) if tok_emb is not None else vis_emb
-                all_embeds.append(seq_emb)
-                if vis_mask is not None:
-                    tok_mask = vis_mask.new_ones(len(ids)) if ids else vis_mask.new_zeros(0)
-                    all_masks.append(torch.cat([vis_mask, tok_mask]) if ids else vis_mask)
-
-            max_len = max(e.shape[0] for e in all_embeds)
-            d_model = all_embeds[0].shape[-1]
-            padded_embeds = vis_emb.new_zeros(num_beams, max_len, d_model)
-            padded_mask   = vis_mask.new_zeros(num_beams, max_len) if vis_mask is not None else None
-            for idx, emb in enumerate(all_embeds):
-                padded_embeds[idx, :emb.shape[0]] = emb
-                if padded_mask is not None:
-                    padded_mask[idx, :emb.shape[0]] = all_masks[idx]
-
-            logits = self._forward_embeds(padded_embeds, padded_mask)
-            candidates = []
-            for beam_idx, (score, ids, done) in enumerate(beams):
-                if done:
-                    candidates.append((score, ids, True))
-                    continue
-                last_pos = vis_len + len(ids) - 1
-                log_p = torch.log_softmax(logits[beam_idx, last_pos, :], dim=-1)
-                if len(ids) == 0 and beam_idx > 0:
-                    log_p = log_p.fill_(-1e9)
-                for lp, tok in zip(*map(lambda t: t.tolist(), log_p.topk(num_beams))):
-                    candidates.append((score + lp, ids + [tok], tok == eos_id))
-            candidates.sort(key=lambda x: -x[0])
-            beams = candidates[:num_beams]
-            if all(done for _, _, done in beams):
-                break
-
-        best_ids = max(beams, key=lambda x: x[0])[1]
-        if not best_ids:
-            return torch.zeros(1, 0, dtype=torch.long, device=device)
-        return torch.tensor(best_ids, dtype=torch.long, device=device).unsqueeze(0)
-
-
-class Nav2TexModel(PreTrainedModel):
-    config_class = Nav2TexConfig
-    base_model_prefix = "model"
-    main_input_name = "pixel_values"
-
-    def __init__(self, config: Nav2TexConfig):
-        super().__init__(config)
-        self.config = config
-        self.visual_encoder = VisualEncoder(
-            NaViT_Encoder(
-                image_size=(config.image_height, config.max_image_width),
-                patch_size=config.patch_size,
-                dim=config.navit_dim,
-                depth=config.navit_depth,
-                heads=config.navit_heads,
-                mlp_dim=config.navit_mlp_dim,
-                dim_head=config.navit_dim_head,
-                dropout=config.navit_dropout,
-                emb_dropout=config.navit_emb_dropout,
-            ),
-            MLPProjector(
-                vision_hidden_size=config.vision_hidden_size,
-                llm_hidden_size=config.llm_hidden_size,
-                intermediate_size=config.projector_intermediate_size,
-            ),
-            max_visual_tokens=config.max_visual_tokens,
-        )
-        self.decoder = CustomDecoder(config)
-        self.post_init()
-
-    def tie_weights(self):
-        if self.config.decoder_weights_tied:
-            self.decoder.tie_weights()
-        else:
-            self.decoder.untie_weights()
-
-    def _init_weights(self, module):
-        return
-
-    @staticmethod
-    def _to_batched_images(pixel_values):
-        if isinstance(pixel_values, list):
-            return pixel_values
-        if isinstance(pixel_values, torch.Tensor):
-            return [[img] for img in pixel_values]
-        raise TypeError(f"Unsupported pixel_values type: {type(pixel_values)}")
-
-    def forward(self, pixel_values, input_ids=None, attention_mask=None, labels=None, **kwargs):
-        batched_images = self._to_batched_images(pixel_values)
-        ve, vm = self.visual_encoder(batched_images)
-        if input_ids is None:
-            return BaseModelOutput(last_hidden_state=ve)
-        te = self.decoder.get_input_embeddings()(input_ids)
-        inputs_embeds = torch.cat([ve, te], dim=1)
-        am = torch.cat([vm.to(dtype=attention_mask.dtype), attention_mask], dim=1)
-        lv = torch.full((labels.shape[0], ve.shape[1]), -100, dtype=labels.dtype, device=labels.device)
-        out = self.decoder(
-            inputs_embeds=inputs_embeds,
-            attention_mask=am,
-            labels=torch.cat([lv, labels], dim=1),
-        )
-        return BaseModelOutput(last_hidden_state=out.last_hidden_state, hidden_states=(out.hidden_states[0],))
-
-    @torch.no_grad()
-    def generate(self, pixel_values, max_new_tokens=None, num_beams=None):
-        batched_images = self._to_batched_images(pixel_values)
-        ve, vm = self.visual_encoder(batched_images)
-        batch = ve.shape[0]
-        bos_id = self.config.decoder_arch["bos_id"]
-        bos_emb = self.decoder.get_input_embeddings()(
-            torch.full((batch, 1), bos_id, dtype=torch.long, device=ve.device)
-        )
-        inputs_embeds = torch.cat([ve, bos_emb], dim=1)
-        attention_mask = torch.cat([
-            vm.to(dtype=torch.long),
-            torch.ones(batch, 1, dtype=torch.long, device=ve.device)
-        ], dim=1)
-        return self.decoder.generate(
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            max_new_tokens=max_new_tokens or self.config.max_new_tokens,
-            num_beams=num_beams or self.config.num_beams,
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from functools import partial
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence as orig_pad_sequence
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+try:
+    from .configuration_latex_decoder import LaTeXDecoderConfig
+    from .configuration_latex_ocr import Nav2TexConfig
+    from .modeling_latex_decoder import LaTeXDecoderForCausalLM
+except ImportError:
+    from nav2tex.configuration_latex_decoder import LaTeXDecoderConfig
+    from nav2tex.configuration_latex_ocr import Nav2TexConfig
+    from nav2tex.modeling_latex_decoder import LaTeXDecoderForCausalLM
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import pad_input, unpad_input
+    HAS_FLASH_ATTN = True
+except ImportError:
+    HAS_FLASH_ATTN = False
+
+
+def exists(val):
+    return val is not None
+
+
+def divisible_by(numer, denom):
+    return (numer % denom) == 0
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.normalized_shape = (dim,)
+        self.eps = 1e-5
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.bias = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        return F.layer_norm(
+            x.float(), self.normalized_shape,
+            self.weight.float(), self.bias.float(), self.eps,
+        ).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, heads, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, 1, dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim=-1) * self.scale * self.gamma.to(x.dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_2d_rope(q, k, h_idx, w_idx):
+    _, _, _, d = q.shape
+    if d % 4 != 0:
+        raise ValueError(f"apply_2d_rope expects dim_head divisible by 4, got D={d}")
+    dim_half = d // 2
+    dim_quarter = d // 4
+    inv_freq = 1.0 / (10000 ** (torch.arange(dim_quarter, device=q.device).float() / dim_quarter))
+    h_theta = h_idx[..., None].float() * inv_freq
+    w_theta = w_idx[..., None].float() * inv_freq
+    sin_h = torch.cat([h_theta.sin(), h_theta.sin()], dim=-1).to(q.dtype)[:, None, :, :]
+    cos_h = torch.cat([h_theta.cos(), h_theta.cos()], dim=-1).to(q.dtype)[:, None, :, :]
+    sin_w = torch.cat([w_theta.sin(), w_theta.sin()], dim=-1).to(q.dtype)[:, None, :, :]
+    cos_w = torch.cat([w_theta.cos(), w_theta.cos()], dim=-1).to(q.dtype)[:, None, :, :]
+
+    def rope(x, sin, cos):
+        return x * cos + rotate_half(x) * sin
+
+    q = torch.cat([rope(q[..., :dim_half], sin_h, cos_h), rope(q[..., dim_half:], sin_w, cos_w)], dim=-1)
+    k = torch.cat([rope(k[..., :dim_half], sin_h, cos_h), rope(k[..., dim_half:], sin_w, cos_w)], dim=-1)
+    return q, k
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.norm = LayerNorm(dim)
+        self.q_norm = RMSNorm(heads, dim_head)
+        self.k_norm = RMSNorm(heads, dim_head)
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim, bias=False), nn.Dropout(dropout))
+
+    def forward(self, x, mask=None, attn_mask=None, positions=None):
+        x = self.norm(x)
+        q = self.to_q(x)
+        k, v = self.to_kv(x).chunk(2, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (q, k, v))
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if positions is not None:
+            q, k = apply_2d_rope(q, k, positions[0], positions[1])
+
+        if HAS_FLASH_ATTN and x.is_cuda and attn_mask is None:
+            fa_dtype = q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
+            q_ = rearrange(q, "b h n d -> b n h d").contiguous().to(fa_dtype)
+            k_ = rearrange(k, "b h n d -> b n h d").contiguous().to(fa_dtype)
+            v_ = rearrange(v, "b h n d -> b n h d").contiguous().to(fa_dtype)
+            if exists(mask):
+                batch, seqlen = mask.shape
+                q_unpad, indices, cu_q, max_q, *_ = unpad_input(q_, mask)
+                k_unpad, _, cu_k, max_k, *_ = unpad_input(k_, mask)
+                v_unpad, _, _, _, *_ = unpad_input(v_, mask)
+                out_unpad = flash_attn_varlen_func(
+                    q_unpad, k_unpad, v_unpad,
+                    cu_seqlens_q=cu_q, cu_seqlens_k=cu_k,
+                    max_seqlen_q=max_q, max_seqlen_k=max_k,
+                    dropout_p=self.dropout.p if self.training else 0.0,
+                    causal=False,
+                )
+                out = pad_input(out_unpad, indices, batch, seqlen)
+            else:
+                out = flash_attn_func(
+                    q_, k_, v_,
+                    dropout_p=self.dropout.p if self.training else 0.0,
+                    causal=False,
+                )
+            out = rearrange(out, "b n h d -> b n (h d)").to(x.dtype)
+        else:
+            dots = torch.matmul(q, k.transpose(-1, -2))
+            if exists(mask):
+                dots = dots.masked_fill(~mask[:, None, None, :], -torch.finfo(dots.dtype).max)
+            if exists(attn_mask):
+                dots = dots.masked_fill(~attn_mask, -torch.finfo(dots.dtype).max)
+            attn = self.dropout(self.attend(dots))
+            out = rearrange(torch.matmul(attn, v), "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.0):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            nn.ModuleList([Attention(dim, heads, dim_head, dropout), FeedForward(dim, mlp_dim, dropout)])
+            for _ in range(depth)
+        ])
+        self.norm = LayerNorm(dim)
+
+    def forward(self, x, mask=None, attn_mask=None, positions=None):
+        for attn, ff in self.layers:
+            x = attn(x, mask=mask, attn_mask=attn_mask, positions=positions) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+
+class NaViT_Encoder(nn.Module):
+    def __init__(self, *, image_size, patch_size, dim, depth, heads, mlp_dim,
+                 channels=3, dim_head=64, dropout=0.0, emb_dropout=0.0):
+        super().__init__()
+        image_height, image_width = image_size
+        assert divisible_by(image_height, patch_size)
+        assert divisible_by(image_width, patch_size)
+        self.patch_size = patch_size
+        self.to_patch_embedding = nn.Sequential(
+            LayerNorm(channels * patch_size ** 2),
+            nn.Linear(channels * patch_size ** 2, dim),
+            LayerNorm(dim),
+        )
+        self.dropout = nn.Dropout(emb_dropout)
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, batched_images):
+        p = self.patch_size
+        device = self.device
+        arange = partial(torch.arange, device=device)
+        pad_sequence = partial(orig_pad_sequence, batch_first=True)
+        batched_sequences, batched_positions = [], []
+
+        for images in batched_images:
+            sequences, positions = [], []
+            for image in images:
+                _, h, w = image.shape
+                ph, pw = h // p, w // p
+                seq = rearrange(image, "c (h p1) (w p2) -> (h w) (c p1 p2)", p1=p, p2=p)
+                pos = torch.stack(torch.meshgrid(arange(ph), arange(pw), indexing="ij"), dim=-1)
+                sequences.append(seq)
+                positions.append(rearrange(pos, "h w c -> (h w) c"))
+            batched_sequences.append(torch.cat(sequences, dim=0))
+            batched_positions.append(torch.cat(positions, dim=0))
+
+        patches = pad_sequence(batched_sequences)
+        patch_positions = pad_sequence(batched_positions)
+        lengths = torch.tensor([seq.shape[0] for seq in batched_sequences], device=device)
+        mask = torch.arange(patches.shape[1], device=device)[None, :] < lengths[:, None]
+        x = self.to_patch_embedding(patches.to(next(self.parameters()).dtype))
+        h_idx, w_idx = patch_positions.unbind(dim=-1)
+        x = self.dropout(x)
+        x = self.transformer(x, mask=mask, positions=(h_idx, w_idx))
+        return x, mask
+
+
+class MLPProjector(nn.Module):
+    def __init__(self, vision_hidden_size=1024, llm_hidden_size=512, intermediate_size=2048):
+        super().__init__()
+        self.norm = nn.LayerNorm(vision_hidden_size)
+        self.gate_proj = nn.Linear(vision_hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(vision_hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, llm_hidden_size, bias=False)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class VisualEncoder(nn.Module):
+    def __init__(self, encoder, bridge, max_visual_tokens):
+        super().__init__()
+        self.navit = encoder
+        self.projector = bridge
+        self.max_visual_tokens = max_visual_tokens
+
+    def forward(self, batched_images):
+        x, mask = self.navit(batched_images)
+        if x.shape[1] > self.max_visual_tokens:
+            x = x[:, :self.max_visual_tokens]
+            mask = mask[:, :self.max_visual_tokens]
+        return self.projector(x), mask
+
+
+class CustomDecoder(nn.Module):
+    def __init__(self, config: Nav2TexConfig):
+        super().__init__()
+        dec = config.decoder_arch
+        self._model = LaTeXDecoderForCausalLM(
+            LaTeXDecoderConfig(
+                vocab_size=dec["vocab_size"],
+                pad_id=dec["pad_id"],
+                bos_id=dec["bos_id"],
+                eos_id=dec["eos_id"],
+                d_model=dec["d_model"],
+                n_heads=dec["n_heads"],
+                n_layers=dec["n_layers"],
+                d_ff=dec["d_ff"],
+                dropout=dec.get("dropout", 0.1),
+                max_seq_len=dec["max_seq_len"],
+                rope_theta=dec.get("rope_theta", 10000.0),
+                tie_weights=dec.get("tie_weights", True),
+            )
+        )
+        self.pad_token_id = self._model.config.pad_id
+        self.eos_token_id = self._model.config.eos_id
+        self._vocab_size = self._model.config.vocab_size
+        self._pad_id = self._model.config.pad_id
+        if not config.decoder_weights_tied:
+            self.untie_weights()
+
+    def get_input_embeddings(self):
+        return self._model.embed_tokens
+
+    def tie_weights(self):
+        self._model.lm_head.weight = self._model.embed_tokens.weight
+
+    def untie_weights(self):
+        if self.are_weights_tied():
+            self._model.lm_head.weight = nn.Parameter(self._model.embed_tokens.weight.detach().clone())
+
+    def are_weights_tied(self):
+        return self._model.lm_head.weight.data_ptr() == self._model.embed_tokens.weight.data_ptr()
+
+    def _forward_embeds(self, inputs_embeds, attention_mask=None):
+        x = self._model.embed_drop(inputs_embeds)
+        mask = attention_mask.bool() if attention_mask is not None else None
+        for layer in self._model.layers:
+            x = layer(x, mask)
+        return self._model.lm_head(self._model.norm_final(x))
+
+    def forward(self, inputs_embeds=None, attention_mask=None, labels=None, **kwargs):
+        logits = self._forward_embeds(inputs_embeds, attention_mask)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1].contiguous()
+            shift_labels = labels[:, 1:].contiguous().masked_fill(
+                labels[:, 1:].contiguous() == self._pad_id, -100
+            )
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self._vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return BaseModelOutput(last_hidden_state=logits, hidden_states=(loss,))
+
+    @torch.no_grad()
+    def generate(self, inputs_embeds, attention_mask, max_new_tokens, num_beams=1):
+        device = inputs_embeds.device
+        batch  = inputs_embeds.shape[0]
+
+        if num_beams > 1:
+            # beam search: only supports batch_size=1
+            assert batch == 1, "beam search only supports batch_size=1"
+            return self._beam_search(inputs_embeds, attention_mask, max_new_tokens, num_beams)
+
+        return self._greedy_batch(inputs_embeds, attention_mask, max_new_tokens)
+
+    @torch.no_grad()
+    def _greedy_batch(self, inputs_embeds, attention_mask, max_new_tokens):
+        """Greedy decoding with true batch support."""
+        eos_id  = self.eos_token_id
+        pad_id  = self._pad_id
+        device  = inputs_embeds.device
+        batch   = inputs_embeds.shape[0]
+        d_model = inputs_embeds.shape[-1]
+
+        # generated token ids per sample, and finished flags
+        gen_ids  = [[] for _ in range(batch)]
+        finished = torch.zeros(batch, dtype=torch.bool, device=device)
+
+        cur_embeds = inputs_embeds                          # (B, vis_len, D)
+        cur_mask   = attention_mask                         # (B, vis_len)
+
+        for _ in range(max_new_tokens):
+            logits   = self._forward_embeds(cur_embeds, cur_mask)  # (B, seq, vocab)
+            next_tok = logits[:, -1, :].argmax(dim=-1)             # (B,)
+
+            for i in range(batch):
+                if not finished[i]:
+                    gen_ids[i].append(next_tok[i].item())
+            finished |= (next_tok == eos_id)
+            if finished.all():
+                break
+
+            tok_emb  = self._model.embed_tokens(next_tok.unsqueeze(1))   # (B, 1, D)
+            tok_mask = cur_mask.new_ones(batch, 1)
+            cur_embeds = torch.cat([cur_embeds, tok_emb], dim=1)
+            cur_mask   = torch.cat([cur_mask,   tok_mask], dim=1)
+
+        # pad to same length and return (B, max_len)
+        max_len = max((len(ids) for ids in gen_ids), default=0)
+        if max_len == 0:
+            return torch.zeros(batch, 0, dtype=torch.long, device=device)
+        out = torch.full((batch, max_len), pad_id, dtype=torch.long, device=device)
+        for i, ids in enumerate(gen_ids):
+            if ids:
+                out[i, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device)
+        return out
+
+    @torch.no_grad()
+    def _beam_search(self, inputs_embeds, attention_mask, max_new_tokens, num_beams):
+        """Original beam search (batch_size=1 only)."""
+        eos_id   = self.eos_token_id
+        device   = inputs_embeds.device
+        vis_emb  = inputs_embeds[0]
+        vis_len  = vis_emb.shape[0]
+        vis_mask = attention_mask[0] if attention_mask is not None else None
+        beams    = [(0.0, [], False) for _ in range(num_beams)]
+
+        for _ in range(max_new_tokens):
+            all_embeds, all_masks = [], []
+            for score, ids, _ in beams:
+                tok_emb = self._model.embed_tokens(torch.tensor(ids, device=device, dtype=torch.long)) if ids else None
+                seq_emb = torch.cat([vis_emb, tok_emb], dim=0) if tok_emb is not None else vis_emb
+                all_embeds.append(seq_emb)
+                if vis_mask is not None:
+                    tok_mask = vis_mask.new_ones(len(ids)) if ids else vis_mask.new_zeros(0)
+                    all_masks.append(torch.cat([vis_mask, tok_mask]) if ids else vis_mask)
+
+            max_len = max(e.shape[0] for e in all_embeds)
+            d_model = all_embeds[0].shape[-1]
+            padded_embeds = vis_emb.new_zeros(num_beams, max_len, d_model)
+            padded_mask   = vis_mask.new_zeros(num_beams, max_len) if vis_mask is not None else None
+            for idx, emb in enumerate(all_embeds):
+                padded_embeds[idx, :emb.shape[0]] = emb
+                if padded_mask is not None:
+                    padded_mask[idx, :emb.shape[0]] = all_masks[idx]
+
+            logits = self._forward_embeds(padded_embeds, padded_mask)
+            candidates = []
+            for beam_idx, (score, ids, done) in enumerate(beams):
+                if done:
+                    candidates.append((score, ids, True))
+                    continue
+                last_pos = vis_len + len(ids) - 1
+                log_p = torch.log_softmax(logits[beam_idx, last_pos, :], dim=-1)
+                if len(ids) == 0 and beam_idx > 0:
+                    log_p = log_p.fill_(-1e9)
+                for lp, tok in zip(*map(lambda t: t.tolist(), log_p.topk(num_beams))):
+                    candidates.append((score + lp, ids + [tok], tok == eos_id))
+            candidates.sort(key=lambda x: -x[0])
+            beams = candidates[:num_beams]
+            if all(done for _, _, done in beams):
+                break
+
+        best_ids = max(beams, key=lambda x: x[0])[1]
+        if not best_ids:
+            return torch.zeros(1, 0, dtype=torch.long, device=device)
+        return torch.tensor(best_ids, dtype=torch.long, device=device).unsqueeze(0)
+
+
+class Nav2TexModel(PreTrainedModel):
+    config_class = Nav2TexConfig
+    base_model_prefix = "model"
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: Nav2TexConfig):
+        super().__init__(config)
+        self.config = config
+        self.visual_encoder = VisualEncoder(
+            NaViT_Encoder(
+                image_size=(config.image_height, config.max_image_width),
+                patch_size=config.patch_size,
+                dim=config.navit_dim,
+                depth=config.navit_depth,
+                heads=config.navit_heads,
+                mlp_dim=config.navit_mlp_dim,
+                dim_head=config.navit_dim_head,
+                dropout=config.navit_dropout,
+                emb_dropout=config.navit_emb_dropout,
+            ),
+            MLPProjector(
+                vision_hidden_size=config.vision_hidden_size,
+                llm_hidden_size=config.llm_hidden_size,
+                intermediate_size=config.projector_intermediate_size,
+            ),
+            max_visual_tokens=config.max_visual_tokens,
+        )
+        self.decoder = CustomDecoder(config)
+        self.post_init()
+
+    def tie_weights(self, **kwargs):
+        if self.config.decoder_weights_tied:
+            self.decoder.tie_weights()
+        else:
+            self.decoder.untie_weights()
+
+    def _init_weights(self, module):
+        return
+
+    @staticmethod
+    def _to_batched_images(pixel_values):
+        if isinstance(pixel_values, list):
+            return pixel_values
+        if isinstance(pixel_values, torch.Tensor):
+            return [[img] for img in pixel_values]
+        raise TypeError(f"Unsupported pixel_values type: {type(pixel_values)}")
+
+    def forward(self, pixel_values, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        batched_images = self._to_batched_images(pixel_values)
+        ve, vm = self.visual_encoder(batched_images)
+        if input_ids is None:
+            return BaseModelOutput(last_hidden_state=ve)
+        te = self.decoder.get_input_embeddings()(input_ids)
+        inputs_embeds = torch.cat([ve, te], dim=1)
+        am = torch.cat([vm.to(dtype=attention_mask.dtype), attention_mask], dim=1)
+        lv = torch.full((labels.shape[0], ve.shape[1]), -100, dtype=labels.dtype, device=labels.device)
+        out = self.decoder(
+            inputs_embeds=inputs_embeds,
+            attention_mask=am,
+            labels=torch.cat([lv, labels], dim=1),
+        )
+        return BaseModelOutput(last_hidden_state=out.last_hidden_state, hidden_states=(out.hidden_states[0],))
+
+    @torch.no_grad()
+    def generate(self, pixel_values, max_new_tokens=None, num_beams=None):
+        batched_images = self._to_batched_images(pixel_values)
+        ve, vm = self.visual_encoder(batched_images)
+        batch = ve.shape[0]
+        bos_id = self.config.decoder_arch["bos_id"]
+        bos_emb = self.decoder.get_input_embeddings()(
+            torch.full((batch, 1), bos_id, dtype=torch.long, device=ve.device)
+        )
+        inputs_embeds = torch.cat([ve, bos_emb], dim=1)
+        attention_mask = torch.cat([
+            vm.to(dtype=torch.long),
+            torch.ones(batch, 1, dtype=torch.long, device=ve.device)
+        ], dim=1)
+        return self.decoder.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens or self.config.max_new_tokens,
+            num_beams=num_beams or self.config.num_beams,
         )