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]