Spaces:
Running on Zero
Running on Zero
| 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] | |