s2sr_pipe/model/architecture.py

"""
sen2sr.model.architecture
=========================
Official SEN2SR model loader (ESAOpenSR / tacofoundation).

This module replaces the custom RRDB implementation with the actual pretrained
SEN2SR model downloaded from HuggingFace via the mlstac library.

Model variants
--------------
Three variants are available on HuggingFace (tacofoundation/sen2sr):

  ┌───────────────────────────────┬────────────┬────────────┬─────────┐
  │ Variant                       │ In bands   │ In / Out   │ Scale   │
  ├───────────────────────────────┼────────────┼────────────┼─────────┤
  │ SEN2SRLite/main  ← default    │ 10 bands   │ (10,H,W)   │   ×4    │
  │ SEN2SRLite/NonReference_RGBN  │ 4 bands    │ (4,H,W)    │   ×4    │
  │ SEN2SRLite/Reference_RSWIR_x2 │ 10 bands   │ (10,H,W)   │   ×2    │
  └───────────────────────────────┴────────────┴────────────┴─────────┘

Default (SEN2SRLite/main):
  - Input  : (B, 10, H, W) float32 [0, 1], bands in order
               B02 B03 B04 B05 B06 B07 B08 B8A B11 B12
  - Output : (B, 10, H×4, W×4) float32
  - Pixel size  : 10 m → 2.5 m
  - Weights     : auto-downloaded from HuggingFace on first call (~100 MB)

The download is idempotent: if the model directory already exists it is reused.
No HuggingFace account or API token is required (public model).
"""
from __future__ import annotations

import os
from pathlib import Path
from typing import Literal

import torch
import torch.nn as nn

from s2sr_pipe.utils.logging_utils import get_logger

logger = get_logger("architecture")

# ── HuggingFace URLs ──────────────────────────────────────────────────────────
_HF_BASE = "https://huggingface.co/tacofoundation/sen2sr/resolve/main"

VARIANT_URLS: dict[str, str] = {
    "main":      f"{_HF_BASE}/SEN2SRLite/main/mlm.json",
    "rgbn_x4":   f"{_HF_BASE}/SEN2SRLite/NonReference_RGBN_x4/mlm.json",
    "rswir_x2":  f"{_HF_BASE}/SEN2SRLite/Reference_RSWIR_x2/mlm.json",
}

VARIANT_SCALE: dict[str, int] = {
    "main":     4,
    "rgbn_x4":  4,
    "rswir_x2": 2,
}

VARIANT_IN_CHANNELS: dict[str, int] = {
    "main":     10,
    "rgbn_x4":  4,
    "rswir_x2": 10,
}

ModelVariant = Literal["main", "rgbn_x4", "rswir_x2"]


# ── Model loader ──────────────────────────────────────────────────────────────

def build_model(
    variant: ModelVariant = "main",
    model_dir: str | Path = "sen2sr_model",
    device: torch.device | str = "cpu",
    force_download: bool = False,
) -> nn.Module:
    """
    Download (once) and load the official SEN2SR pretrained model.

    Parameters
    ----------
    variant       : Which SEN2SR variant to use (see module docstring).
    model_dir     : Local directory where the model files are stored.
                    The directory is created if it doesn't exist.
                    Re-using the same directory avoids re-downloading.
    device        : Target device for inference.
    force_download: Re-download even if model_dir already exists.

    Returns
    -------
    torch.nn.Module in eval mode, on *device*.

    Raises
    ------
    ImportError  : if `mlstac` is not installed.
    ValueError   : if an unknown variant name is given.
    RuntimeError : if the download or model compilation fails.

    Notes
    -----
    * mlstac.download() is idempotent: it skips the download if the target
      directory already contains a valid model.
    * compiled_model() returns a standard torch.nn.Module (likely TorchScript
      or a wrapped nn.Module), fully compatible with torch.no_grad() and AMP.
    * The model expects patches of exactly 128×128 pixels at the LR side.
      Larger inputs are handled by sen2sr.predict_large() in inference.py.
    """
    try:
        import mlstac
    except ImportError as exc:
        raise ImportError(
            "mlstac is required to load the official SEN2SR model.\n"
            "Install it with:  pip install mlstac sen2sr"
        ) from exc

    if variant not in VARIANT_URLS:
        raise ValueError(
            f"Unknown variant '{variant}'. "
            f"Choose from: {list(VARIANT_URLS.keys())}"
        )

    model_dir = Path(model_dir).resolve()
    mlm_json = model_dir / "mlm.json"

    # ── Download ──────────────────────────────────────────────────────────────
    if force_download or not mlm_json.exists():
        logger.info(
            "Downloading SEN2SR weights (%s) from HuggingFace -> %s",
            variant,
            model_dir,
        )
        logger.info("URL: %s", VARIANT_URLS[variant])
        mlstac.download(
            file=VARIANT_URLS[variant],
            output_dir=str(model_dir),
        )
        logger.info("Download complete.")
    else:
        logger.info(
            "Model already present at %s - skipping download. "
            "Use force_download=True to re-download.",
            model_dir,
        )

    # ── Load & compile ────────────────────────────────────────────────────────
    logger.info("Loading compiled model from %s ...", model_dir)
    container = mlstac.load(str(model_dir))
    model: nn.Module = container.compiled_model(device=device)
    model = model.to(device)
    model.eval()

    scale = VARIANT_SCALE[variant]
    in_ch = VARIANT_IN_CHANNELS[variant]
    logger.info(
        "SEN2SR model ready | variant=%s | in_channels=%d | scale=x%d | device=%s",
        variant,
        in_ch,
        scale,
        device,
    )
    return model


def get_scale(variant: ModelVariant = "main") -> int:
    """Return the upscale factor for a given variant."""
    return VARIANT_SCALE[variant]


def get_in_channels(variant: ModelVariant = "main") -> int:
    """Return the number of input channels expected by a given variant."""
    return VARIANT_IN_CHANNELS[variant]