core/training_utils.py

"""Shared trainer utilities: EMA, config/paths, DINO/fusion builders and argparse helpers."""

import os
from datetime import datetime

import torch
import yaml

from models.hyperfeature_fusion import create_hyperfeature_fusion
from models.dinov3_hf_extractor import create_dinov3_hf_extractor


DINO_REPO_MAP = {
    "dinov3_vits16": "facebook/dinov3-vits16-pretrain-lvd1689m",
    "dinov3_vitb16": "facebook/dinov3-vitb16-pretrain-lvd1689m",
    "dinov3_vitl16": "facebook/dinov3-vitl16-pretrain-lvd1689m",
    "dinov3_vitg16": "facebook/dinov3-vitg16-pretrain-lvd1689m",
}


class EMA:
    """Exponential Moving Average of trainable, floating-point model parameters."""

    def __init__(self, model, beta=0.999):
        self.beta = beta
        self.shadow = {}
        self.backup = {}
        for name, param in model.named_parameters():
            if param.requires_grad and param.dtype.is_floating_point:
                self.shadow[name] = param.data.detach().clone()

    def update(self, model):
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in self.shadow:
                    if self.shadow[name].device != param.device:
                        self.shadow[name] = self.shadow[name].to(param.device)
                    self.shadow[name].mul_(self.beta).add_(param.data.detach(), alpha=1.0 - self.beta)

    def apply_shadow(self, model):
        self.backup = {}
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in self.shadow:
                    if self.shadow[name].device != param.device:
                        self.shadow[name] = self.shadow[name].to(param.device)
                    self.backup[name] = param.data.detach().clone()
                    param.data.copy_(self.shadow[name])
        return self.backup

    def restore_backup(self, model):
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in self.backup:
                    param.data.copy_(self.backup[name])
        self.backup = {}

    def restore(self, model, backup=None):
        b = backup if backup is not None else self.backup
        with torch.no_grad():
            for name, param in model.named_parameters():
                if name in b:
                    if b[name].device != param.device:
                        b[name] = b[name].to(param.device)
                    param.data.copy_(b[name])
        self.backup = {}


def _expand_env(value):
    """Recursively expand ${VAR}/$VAR and ~ in all string values of a config."""
    if isinstance(value, str):
        return os.path.expanduser(os.path.expandvars(value))
    if isinstance(value, dict):
        return {k: _expand_env(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_expand_env(v) for v in value]
    return value


def load_config(config_path: str):
    with open(config_path, "r") as f:
        return _expand_env(yaml.safe_load(f))


def setup_paths(config, task: str):
    """Timestamped log/checkpoint dirs plus the (permanent) feature cache dir."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    cache_dir = config["paths"].get("permanent_cache_dir") or f"{config['paths']['cache_base_dir']}/{task}_cache"
    return {
        "cache_dir": cache_dir,
        "log_dir": f"{config['paths']['log_base_dir']}/{task}_logs_{timestamp}",
        "checkpoint_dir": f"{config['paths']['checkpoint_base_dir']}/{task}_checkpoints_{timestamp}",
    }


def calculate_distributed_concept_channels(num_concepts, num_features=4):
    """Per-layer concept channel counts when distributing concepts across layers."""
    per = num_concepts // num_features
    rem = num_concepts % num_features
    return [per + (1 if i < rem else 0) for i in range(num_features)]


def feature_mode_flags(feature_mode):
    """Return ``(use_flux, use_dino)`` for an ablation mode."""
    assert feature_mode in ("full", "flux_only", "dino_only"), feature_mode
    return feature_mode in ("full", "flux_only"), feature_mode in ("full", "dino_only")


def resolve_c_dino(feature_mode, dino_model):
    """DINO feature width for the given backbone (0 when DINO is disabled)."""
    if feature_mode == "flux_only":
        return 0
    if "vits" in dino_model:
        return 384
    if "vitl" in dino_model:
        return 1024
    if "vitg" in dino_model:
        return 1536
    return 768  # vitb


def build_dino_extractor(dino_model, feature_mode, layer_indices=(2, 5, 8, 11)):
    """Build the DINOv3 extractor (HF when known, torch.hub fallback otherwise)."""
    trainable = feature_mode == "dino_only"
    repo_id = DINO_REPO_MAP.get(dino_model)
    if repo_id:
        return create_dinov3_hf_extractor(
            repo_id=repo_id, take_indices=list(layer_indices), trainable=trainable
        )
    from models.dino_fusion import create_dino_extractor
    return create_dino_extractor(model_name=dino_model, take_indices=list(layer_indices))


def build_hyperfeature_fusion(use_flux, num_timesteps, hidden_dim, num_transformer_layers,
                              layer_scale_init, fusion_type="transformer", return_alpha=False):
    """Build the multi-timestep hyperfeature fusion module, or ``None`` without FLUX."""
    if not use_flux:
        return None
    return create_hyperfeature_fusion(
        num_timesteps=num_timesteps,
        num_layers=4,
        fusion_type=fusion_type,
        hidden_dim=hidden_dim,
        num_transformer_layers=num_transformer_layers,
        layer_scale_init=layer_scale_init,
        return_alpha=return_alpha,
    )


def add_shared_training_args(parser, *, config_default, hidden_dim, num_transformer_layers,
                             layer_scale_init, include_cache_args=True):
    """Register the backbone/ablation arguments common to all task trainers."""
    parser.add_argument("--config", type=str, default=config_default)
    parser.add_argument("--num_timesteps", type=int, default=4)
    parser.add_argument("--hidden_dim", type=int, default=hidden_dim)
    parser.add_argument("--num_transformer_layers", type=int, default=num_transformer_layers)
    parser.add_argument("--layer_scale_init", type=float, default=layer_scale_init)
    parser.add_argument("--dino_model", type=str, default="dinov3_vitb16")
    parser.add_argument("--feature_mode", type=str, default="full",
                        choices=["full", "flux_only", "dino_only"],
                        help="Ablation: 'full' (FLUX+concepts+DINO), 'flux_only' (FLUX+concepts), "
                             "'dino_only' (trainable DINO only)")
    parser.add_argument("--fast_dev_run", action="store_true",
                        help="Smoke test: run a single train+val batch then exit (verifies the "
                             "training pipeline starts end-to-end)")
    if include_cache_args:
        parser.add_argument("--limit_images", type=int, default=None, help="Limit dataset size for testing")
        parser.add_argument("--cache_only", action="store_true",
                            help="Load features from cache only (no FLUX pipeline)")


def add_model_args(parser):
    """Architecture flags for inference/generation; must match the checkpoint."""
    parser.add_argument("--dino_model", default="dinov3_vitb16")
    parser.add_argument("--num_timesteps", type=int, default=4)
    parser.add_argument("--hidden_dim", type=int, default=768)
    parser.add_argument("--num_transformer_layers", type=int, default=3)
    parser.add_argument("--layer_scale_init", type=float, default=1e-6)
    parser.add_argument("--fusion_type", default="transformer", choices=["attention", "transformer"])
    parser.add_argument("--high_res_decoder", action="store_true", help="(nyu) high-res depth decoder")
    parser.add_argument("--device", default="cuda")