model_loader.py

"""
model_loader.py

Singleton loader untuk Visual Encoder (PlantField Image Search).
Model di-download otomatis dari Hugging Face Hub saat startup.

Arsitektur:
    Visual Encoder : Swin-T (28M params) + Projection Head
    Shared Space   : 512-dim, L2-normalized
    Total          : ~28M params (visual encoder only)

Hanya visual encoder yang di-load untuk image embedding.
Text encoder dan classifier diabaikan (partial load).
"""

import os
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from huggingface_hub import hf_hub_download

# ── Config ────────────────────────────────────────────────
MODEL_REPO_ID   = os.getenv("MODEL_REPO_ID", "deadbear34/plantfield-scold")
MODEL_FILENAME  = os.getenv("MODEL_FILENAME", "scold_plantfield_final.pt")
HF_TOKEN        = os.getenv("HF_TOKEN", None)


# ── Visual Encoder ────────────────────────────────────────

class VisualEncoder(nn.Module):
    """Swin-T backbone + MLP projection head."""

    def __init__(self, embed_dim: int = 512, dropout: float = 0.2):
        super().__init__()
        self.backbone = timm.create_model(
            "swin_tiny_patch4_window7_224",
            pretrained=False,
            num_classes=0,
            global_pool="avg",
        )
        D = self.backbone.num_features  # 768
        self.proj = nn.Sequential(
            nn.LayerNorm(D),
            nn.Linear(D, D),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(D, embed_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(self.proj(self.backbone(x)), dim=-1)


class ImageEmbedder(nn.Module):
    """
    Image embedding model — extracts visual encoder from SCOLD checkpoint.
    Produces 512-dim L2-normalized embeddings for image similarity search.
    """

    def __init__(self, embed_dim: int = 512):
        super().__init__()
        self.visual_encoder = VisualEncoder(embed_dim)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode image tensor → 512-dim L2-normalized embedding."""
        return self.visual_encoder(x)


# ── Singleton state ───────────────────────────────────────
_model  : Optional[ImageEmbedder] = None
_device : torch.device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)
_embed_dim: int = 512


def load_model() -> ImageEmbedder:
    """
    Download dan load visual encoder dari SCOLD checkpoint di HF Hub.
    Hanya visual encoder yang di-load; text encoder & classifier di-skip.
    Singleton — hanya dipanggil sekali saat startup.
    """
    global _model, _embed_dim

    if _model is not None:
        return _model

    print(f"[PlantField] Loading ImageEmbedder on {_device}...")
    print(f"[PlantField] Downloading from: {MODEL_REPO_ID}/{MODEL_FILENAME}")

    # ── Download checkpoint dari HF Hub ───────────────────
    ckpt_path = hf_hub_download(
        repo_id  = MODEL_REPO_ID,
        filename = MODEL_FILENAME,
        token    = HF_TOKEN if HF_TOKEN else None,
    )
    print(f"[PlantField] Checkpoint downloaded: {ckpt_path}")

    ckpt = torch.load(ckpt_path, map_location=_device, weights_only=False)

    # ── Extract config ────────────────────────────────────
    _embed_dim = ckpt["cfg"]["embed_dim"]

    # ── Inisialisasi model (visual encoder only) ──────────
    _model = ImageEmbedder(embed_dim=_embed_dim)

    # ── Partial load: hanya visual_encoder weights ────────
    full_state = ckpt["model_state"]
    visual_state = {}
    for key, value in full_state.items():
        if key.startswith("visual_encoder."):
            visual_state[key] = value

    _model.load_state_dict(visual_state)
    _model.eval()
    _model.to(_device)

    # ── Summary ───────────────────────────────────────────
    total_params = sum(p.numel() for p in _model.parameters())
    print(f"[PlantField] ✓ ImageEmbedder loaded")
    print(f"  Embed dim  : {_embed_dim}")
    print(f"  Parameters : {total_params / 1e6:.1f}M")
    print(f"  Device     : {_device}")

    return _model


# ── Getters ───────────────────────────────────────────────

def get_model() -> Optional[ImageEmbedder]:
    return _model

def get_device() -> torch.device:
    return _device

def get_embed_dim() -> int:
    return _embed_dim

def is_model_ready() -> bool:
    return _model is not None