ZeroGPU test

app_function.py

@@ -5,6 +5,7 @@ import torch
 import yaml
 import gradio as gr
 import tempfile
+import spaces
 
 # os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 import time
@@ -219,6 +220,7 @@ class YOND_Backend:
             gr.Error(f"图像更新失败: {str(e)}")
             raise RuntimeError(f"图像更新失败: {str(e)}") from e
 
+    @spaces.GPU
     def estimate_noise(self, double_est, ransac, patch_size):
         """执行噪声估计"""
         if not self.yond:
@@ -298,6 +300,7 @@ class YOND_Backend:
             gr.Error(f"噪声估计失败: {str(e)}")
             raise RuntimeError(f"噪声估计失败: {str(e)}") from e
 
+    @spaces.GPU
     def enhance_image(self, gain, sigma, sigsnr, ddim_mode, patch_size):
         """执行图像增强"""
         if not self.yond:
@@ -440,7 +443,7 @@ class YOND_Backend:
         return tmp_file_path_png
 
 class YONDParser():
-    def __init__(self, yaml_path="runfiles/Gaussian/default_gru32n.yml"):
+    def __init__(self, yaml_path="runfiles/Gaussian/gru32n_paper_noclip.yml"):
         self.runfile = yaml_path
         self.mode = 'eval'
         self.debug = False

requirements.txt

@@ -19,4 +19,5 @@ Pillow
 einops
 timm
 pyarmor
-pytorch_wavelets
+pytorch_wavelets
+spaces