model_loader.py

import os
import time
import torch
import numpy as np
from PIL import Image


def detect_device():
    """Auto-detect device. Returns CPU on ZeroGPU (GPU not available at startup).

    Always uses fp32 — DAT2 architecture produces NaN in fp16 (attention softmax
    + LayerNorm overflow). 48GB VRAM is more than sufficient for fp32.
    """
    if torch.cuda.is_available():
        return "cuda", torch.float32
    return "cpu", torch.float32


# ---------------------------------------------------------------------------
# HuggingFace Hub 模型下载
# ---------------------------------------------------------------------------
# 默认从 Kim2091/UltraSharpV2 公开仓库下载，可通过环境变量覆盖:
#   MODEL_REPO_ID    - 覆盖默认仓库 ID
#   MODEL_FILENAME   - 覆盖默认文件名
#   HF_ENDPOINT      - 镜像站地址，如 https://hf-mirror.com（国内加速）
#   HF_TOKEN         - 私有仓库的 token（公开仓库无需设置）
# ---------------------------------------------------------------------------

_DEFAULT_REPO_ID = "Kim2091/UltraSharpV2"
_DEFAULT_FILENAME = "4x-UltraSharpV2.pth"

MODEL_CANDIDATES = [
    "4x-UltraSharpV2.pth",
    "4x-UltraSharpV2.safetensors",
    "4x-UltraSharpV2.pt",
]


def _download_from_hub(repo_id: str, filename: str) -> str:
    """从 HuggingFace Hub 下载模型文件（自动缓存，重复调用不重新下载）。"""
    from huggingface_hub import hf_hub_download

    token = os.environ.get("HF_TOKEN")
    endpoint = os.environ.get("HF_ENDPOINT")
    if endpoint:
        print(f"[UltraSharpV2] 使用镜像: {endpoint}")

    print(f"[UltraSharpV2] 从 HF Hub 下载: {repo_id}/{filename}")
    path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        token=token,
        endpoint=endpoint,
    )
    print(f"[UltraSharpV2] 下载成功: {path}")
    return path


def _resolve_model_path() -> str:
    """按优先级解析模型路径: 本地文件 > HF Hub（默认 Kim2091/UltraSharpV2）。"""
    # 1) 本地文件优先（存在则直接使用，跳过网络）
    for name in MODEL_CANDIDATES:
        if os.path.exists(name):
            print(f"[UltraSharpV2] 使用本地模型: {name}")
            return name

    # 2) 从 HF Hub 下载
    repo_id = os.environ.get("MODEL_REPO_ID", _DEFAULT_REPO_ID)
    filename = os.environ.get("MODEL_FILENAME", _DEFAULT_FILENAME)
    return _download_from_hub(repo_id, filename)


class UltraSharpV2:
    def __init__(self, model_path=None, device=None):
        """
        Args:
            model_path: path to model file (auto-resolve from HF Hub or local if None).
            device: "cpu" (ZeroGPU default), "cuda", or None (auto-detect).
        """
        if model_path is None:
            model_path = _resolve_model_path()

        if device is not None:
            self.device = device
            self.dtype = torch.float32
        else:
            self.device, self.dtype = detect_device()

        self._model_path = model_path
        self.model = self._load_model(model_path)
        self.scale = self.model.scale
        print(f"[UltraSharpV2] 设备: {self.device}, 精度: {self.dtype}")
        print(f"[UltraSharpV2] 模型加载完毕, 放大倍率: {self.scale}x")

    def _load_model(self, path):
        from spandrel import ModelLoader

        loader = ModelLoader()
        model = loader.load_from_file(path)
        model.model.to(self.device).to(self.dtype).eval()
        return model

    def to_cuda(self):
        """Move model to CUDA (called inside @spaces.GPU decorated function).

        Keeps fp32 — DAT2 architecture produces NaN in fp16.
        """
        if self.device == "cuda":
            return
        print("[UltraSharpV2] 正在将模型移至 GPU ...")
        self.device = "cuda"
        self.model.model.to(self.device)
        torch.cuda.empty_cache()

    def to_cpu(self):
        """Move model back to CPU to release ZeroGPU memory."""
        if self.device == "cpu":
            return
        print("[UltraSharpV2] 正在将模型移回 CPU ...")
        self.model.model.to("cpu")
        self.device = "cpu"
        torch.cuda.empty_cache()

    def upscale(
        self,
        image: Image.Image,
        tile_size: int = 512,
        tile_overlap: int = 32,
        target_scale: float = 4.0,
    ) -> tuple[Image.Image, float]:
        start = time.time()

        tensor = self._pil_to_tensor(image)
        _, _, h, w = tensor.shape

        if h <= tile_size and w <= tile_size:
            with torch.no_grad():
                output = self.model(tensor.to(self.dtype)).float()
        else:
            output = self._tiled_upscale(tensor, tile_size, tile_overlap)

        result = self._tensor_to_pil(output)

        if target_scale > 0 and abs(target_scale - self.scale) > 0.01:
            dest_w = int(w * target_scale)
            dest_h = int(h * target_scale)
            result = result.resize((dest_w, dest_h), Image.LANCZOS)

        elapsed = time.time() - start
        print(
            f"[UltraSharpV2] 推理完成, 尺寸: {h}x{w} -> {result.width}x{result.height}, 耗时: {elapsed:.2f}s"
        )
        return result, elapsed

    def _pil_to_tensor(self, img: Image.Image) -> torch.Tensor:
        if img.mode != "RGB":
            img = img.convert("RGB")
        arr = np.array(img).astype(np.float32) / 255.0
        tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
        return tensor.to(self.device)

    def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
        tensor = tensor.squeeze(0).float()
        if torch.isnan(tensor).any() or torch.isinf(tensor).any():
            raise RuntimeError(
                "模型输出包含 NaN/Inf — 请检查是否使用了 fp16（DAT2 架构不支持 fp16）"
            )
        tensor = tensor.clamp(0, 1)
        arr = (tensor.permute(1, 2, 0).cpu().numpy() * 255).round().astype(np.uint8)
        return Image.fromarray(arr)

    def _tiled_upscale(
        self, tensor: torch.Tensor, tile_size: int, tile_overlap: int
    ) -> torch.Tensor:
        _, c, h, w = tensor.shape
        scale = self.scale
        pad = min(tile_overlap, tile_size // 4)

        padded = torch.nn.functional.pad(
            tensor, (pad, pad, pad, pad), mode="reflect"
        )
        _, _, hp, wp = padded.shape

        out_tile = tile_size * scale
        out_hp = hp * scale
        out_wp = wp * scale
        stride = tile_size - pad * 2

        output = torch.zeros(1, c, out_hp, out_wp, device=self.device, dtype=torch.float32)
        weight = torch.zeros(1, 1, out_hp, out_wp, device=self.device, dtype=torch.float32)

        wy = torch.ones(out_tile, device=self.device)
        wx = torch.ones(out_tile, device=self.device)
        if pad > 0:
            ramp = torch.linspace(0, 1, pad * scale, device=self.device)
            wy[: pad * scale] = ramp
            wy[-pad * scale :] = ramp.flip(0)
            wx[: pad * scale] = ramp
            wx[-pad * scale :] = ramp.flip(0)
        wmap = wy.view(1, 1, -1, 1) * wx.view(1, 1, 1, -1)

        for y in range(0, hp, stride):
            for x in range(0, wp, stride):
                y1 = min(y + tile_size, hp)
                x1 = min(x + tile_size, wp)
                y0 = max(0, y1 - tile_size)
                x0 = max(0, x1 - tile_size)

                tile = padded[:, :, y0:y1, x0:x1]
                with torch.no_grad():
                    out = self.model(tile.to(self.dtype)).float()

                oh, ow = out.shape[2], out.shape[3]
                oy0, ox0 = y0 * scale, x0 * scale
                wc = wmap[:, :, :oh, :ow]

                output[:, :, oy0 : oy0 + oh, ox0 : ox0 + ow] += out * wc
                weight[:, :, oy0 : oy0 + oh, ox0 : ox0 + ow] += wc

        output /= weight.clamp(min=1e-8)

        crop_start = pad * scale
        crop_end_h = crop_start + h * scale
        crop_end_w = crop_start + w * scale
        return output[:, :, crop_start:crop_end_h, crop_start:crop_end_w]