models/dinov2_encoder.py

"""
DINOv2 Encoder — frozen, extracts semantic features from LR images for DiT Cross-Attention.

Uses HF transformers AutoModel for DINOv2 (avoids torch.hub dependency on GitHub).
Uses output_hidden_states=True to get intermediate layer features.
"""

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms import Normalize

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

DINOV2_HF_NAMES = {
    "dinov2b": "facebook/dinov2-base",
    "dinov2l": "facebook/dinov2-large",
    "dinov2g": "facebook/dinov2-giant",
}


class Dinov2Encoder(nn.Module):
    """Frozen DINOv2 feature extractor, outputting features from specified intermediate layers."""

    def __init__(
        self,
        enc_type: str = "dinov2b",
        dinov2_size: int = 448,
        layer_indices: list[int] | None = None,
        device: str = "cuda",
    ):
        super().__init__()

        self.dinov2_size = dinov2_size
        self.layer_indices = layer_indices or [8]

        hf_name = DINOV2_HF_NAMES.get(enc_type)
        if hf_name is None:
            raise ValueError(
                f"Unknown DINOv2 type: {enc_type}, "
                f"expected one of {list(DINOV2_HF_NAMES)}"
            )

        print(f"Loading DINOv2 from HF transformers: {hf_name} ...")
        from transformers import AutoModel
        encoder = AutoModel.from_pretrained(hf_name)

        self.encoder = encoder.to(device).eval()
        for p in self.encoder.parameters():
            p.requires_grad_(False)

        print(f"DINOv2 encoder loaded, layers={self.layer_indices}")

    def preprocess(self, lr: torch.Tensor) -> torch.Tensor:
        """
        lr: [B, 3, H, W] float [0, 1]
        -> resize -> clamp -> ImageNet normalization
        """
        x = F.interpolate(lr, size=self.dinov2_size, mode="bicubic", align_corners=False)
        x = x.clamp(0, 1)
        x = Normalize(IMAGENET_MEAN, IMAGENET_STD)(x)
        return x

    @torch.no_grad()
    def forward(self, lr: torch.Tensor) -> list[torch.Tensor]:
        """
        lr: [B, 3, H, W] float [0, 1]
        -> list of [B, N_patches, enc_dim]
        """
        x = self.preprocess(lr)

        # Use output_hidden_states=True to get all layer outputs
        outputs = self.encoder(x, output_hidden_states=True)
        # hidden_states is a tuple of (num_layers + 1) tensors:
        #   hidden_states[0] = embeddings output
        #   hidden_states[1..N] = output of each encoder layer
        hidden_states = outputs.hidden_states  # tuple of (B, seq_len, hidden_size)

        # Build feature list: layer_i -> hidden_states[i+1] (skip embeddings output)
        # The original code extracts features from intermediate layers (without CLS token)
        z = []
        for idx in self.layer_indices:
            # hidden_states[idx+1] because hidden_states[0] is the embedding output
            # and hidden_states[1..N] are layer outputs
            feat = hidden_states[idx + 1][:, 1:]  # Remove CLS token
            z.append(feat)

        # Replace last with the final norm output (last_hidden_state)
        # The original code does: z[-1] = x_norm (the final normalized output without CLS)
        z[-1] = outputs.last_hidden_state[:, 1:]

        return z


def create_dinov2_encoder(config_path: str, device: str = "cuda") -> Dinov2Encoder | None:
    """Create DINOv2 encoder from YAML config, returns None if not configured."""
    import yaml
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    dv2 = cfg.get("dinov2", {}) or {}
    if not dv2:
        return None

    return Dinov2Encoder(
        enc_type=dv2.get("enc_type", "dinov2b"),
        dinov2_size=dv2.get("dinov2_size", 448),
        layer_indices=dv2.get("layer_dinov2b_list", [8]),
        device=device,
    )