fix fp16 precision bug

model_loader.py

@@ -6,9 +6,13 @@ from PIL import Image
 
 
 def detect_device():
-    """Auto-detect device. Returns CPU on ZeroGPU (GPU not available at startup)."""
+    """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.float16
+        return "cuda", torch.float32
     return "cpu", torch.float32
 
 
@@ -78,7 +82,7 @@ class UltraSharpV2:
 
         if device is not None:
             self.device = device
-            self.dtype = torch.float16 if device == "cuda" else torch.float32
+            self.dtype = torch.float32
         else:
             self.device, self.dtype = detect_device()
 
@@ -97,13 +101,15 @@ class UltraSharpV2:
         return model
 
     def to_cuda(self):
-        """Move model to CUDA (called inside @spaces.GPU decorated function)."""
+        """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.dtype = torch.float16
-        self.model.model.to(self.device).to(self.dtype)
+        self.model.model.to(self.device)
         torch.cuda.empty_cache()
 
     def to_cpu(self):
@@ -111,9 +117,8 @@ class UltraSharpV2:
         if self.device == "cpu":
             return
         print("[UltraSharpV2] 正在将模型移回 CPU ...")
-        self.model.model.to("cpu").to(torch.float32)
+        self.model.model.to("cpu")
         self.device = "cpu"
-        self.dtype = torch.float32
         torch.cuda.empty_cache()
 
     def upscale(
@@ -155,7 +160,12 @@ class UltraSharpV2:
         return tensor.to(self.device)
 
     def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
-        tensor = tensor.squeeze(0).float().clamp(0, 1)
+        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)