ZeroGPU optimize

app.py

@@ -130,7 +130,7 @@ with gr.Blocks(title="YOND WebUI", css="""
                     bl = gr.Number(label="Black Level", value=64.0, precision=1, scale=1)
                     wp = gr.Number(label="White Point", value=1023.0, precision=1, scale=1)
                     ratio = gr.Number(label="DGain (x Ratio)", value=1.0, precision=1, scale=1)
-                    ispgain = gr.Number(label="ISPGain", value=1.0, precision=1, scale=1)
+                    ispgain = gr.Number(label="ISPGain (Visual Only)", value=1.0, precision=1, scale=1)
                 with gr.Row():
                     image_btn = gr.Button("Load Image", variant="primary")
                     image_update_btn = gr.Button("Update Image", variant="secondary")
@@ -144,7 +144,7 @@ with gr.Blocks(title="YOND WebUI", css="""
             
                 est_btn = gr.Button("Noise Estimation", variant="primary")
                 
-                use_ddim = gr.Checkbox(label="DDIM Mode", value=False, scale=1)
+                use_ddim = gr.Checkbox(label="DDIM Mode (Please use large model, e.g., gru64n)", value=False, scale=1)
                 with gr.Row():
                     patch_size = gr.Number(label="Patch Size", value=1024, precision=0)
                     sigsnr = gr.Number(label="SigSNR", precision=2, value=1.03)

app_function.py

@@ -221,6 +221,40 @@ class YOND_Backend:
             raise RuntimeError(f"图像更新失败: {str(e)}") from e
 
     @spaces.GPU
+    def denoise(self, raw_vst, patch_size, nsr):
+        ################# 准备去噪 #################
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.yond.net = self.yond.net.to(self.device)
+        raw_vst = torch.from_numpy(raw_vst).float().to(self.device).permute(2,0,1)[None,]
+        if 'guided' in self.yond.arch:
+            t = torch.tensor(nsr*self.p['sigsnr'], dtype=raw_vst.dtype, device=self.device).view(-1,1,1,1)
+
+            # Denoise & pad
+            target_size = patch_size  # 设置目标块大小，可以根据需要调整 GPU:1024
+            overlap_ratio = 1/8  # 设置重叠率，可以根据需要调整
+            
+            # 使用改进的 big_image_split 函数
+            raw_inp, metadata = big_image_split(raw_vst, target_size, overlap_ratio)
+
+            raw_dn = torch.zeros_like(raw_inp[:,:4])
+            with torch.no_grad():
+                if self.p['ddim_mode']:
+                    for i in range(raw_inp.shape[0]):  # 处理所有切块
+                        print(f'Patch: {i+1}/{len(raw_dn)}')
+                        raw_dn[i] = ddim(raw_inp[i][None,].clip(None, 2), self.yond.net, t, epoch=self.p['epoch'], 
+                                            sigma_t=self.p['sigma_t'], eta=self.p['eta_t'], sigma_corr=1.00)
+                else:
+                    for i in range(raw_inp.shape[0]):  # 处理所有切块
+                        input_tensor = raw_inp[i][None,].clip(None, 2)
+                        raw_dn[i] = self.yond.net(input_tensor, t).clamp(0,None)
+
+            # 使用改进的 big_image_merge 函数
+            raw_dn = big_image_merge(raw_dn, metadata, blend_mode='avg')
+        
+        ################# VST逆变换 #################
+        raw_dn = raw_dn[0].permute(1,2,0).detach().cpu().numpy()
+        return raw_dn
+
     def estimate_noise(self, double_est, ransac, patch_size):
         """执行噪声估计"""
         if not self.yond:
@@ -246,49 +280,30 @@ class YOND_Backend:
 
             if double_est:
                 log(" 使用精估计")
-                lr_raw_np = lr_raw * self.p['scale']
-                ######## EM-VST矫正VST噪图期望偏差 ########
-                bias_base = np.maximum(lr_raw_np, 0)
-                bias = self.bias_lut.get_lut(bias_base, K=self.p['gain'], sigGs=self.p['sigma'])
-                raw_vst = VST(lr_raw_np, self.p['sigma'], gain=self.p['gain'])
-                raw_vst = raw_vst - bias
-
-                ################# VST变换 #################
-                lower = VST(0, self.p['sigma'], gain=self.p['gain'])
-                upper = VST(self.p['scale'], self.p['sigma'], gain=self.p['gain'])
-                nsr = 1 / (upper - lower)
-                raw_vst = (raw_vst - lower) / (upper - lower)
-
-                ################# 准备去噪 #################
-                self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-                self.yond.net = self.yond.net.to(self.device)
-                raw_vst = torch.from_numpy(raw_vst).float().to(self.device).permute(2,0,1)[None,]
-                if 'guided' in self.yond.arch:
-                    sigma_corr = 1.03
-                    t = torch.tensor(nsr*sigma_corr, dtype=raw_vst.dtype, device=self.device)
-
-                    # Denoise & pad
-                    target_size = patch_size  # 设置目标块大小，可以根据需要调整
-                    overlap_ratio = 1/8  # 设置重叠率，可以根据需要调整
-                    
-                    # 使用改进的 big_image_split 函数
-                    raw_inp, metadata = big_image_split(raw_vst, target_size, overlap_ratio)
-
-                    raw_dn = torch.zeros_like(raw_inp[:,:4])
-                    with torch.no_grad():
-                        for i in range(raw_inp.shape[0]):  # 处理所有切块
-                            input_tensor = raw_inp[i][None,].clip(None, 2).to(self.device)
-                            raw_dn[i] = self.yond.net(input_tensor, t).clamp(0,None)
-
-                    # 使用改进的 big_image_merge 函数
-                    raw_dn = big_image_merge(raw_dn, metadata, blend_mode='avg')
-                    
-                ################# VST逆变换 #################
-                raw_vst = raw_dn[0].permute(1,2,0).detach().cpu().numpy()
-                raw_vst = raw_vst * (upper - lower) + lower
-                raw_vst = inverse_VST(raw_vst, self.p['sigma'], gain=self.p['gain']) / self.p['scale']
-
-                reg, self.mask_data = SimpleNLF(rggb2bayer(lr_raw), rggb2bayer(raw_vst), k=13, 
+                if self.denoised_npy is None:
+                    log(" 之前没去噪，先去噪再估计")
+                    lr_raw_np = lr_raw * self.p['scale']
+                    ######## EM-VST矫正VST噪图期望偏差 ########
+                    bias_base = np.maximum(lr_raw_np, 0)
+                    bias = self.bias_lut.get_lut(bias_base, K=self.p['gain'], sigGs=self.p['sigma'])
+                    raw_vst = VST(lr_raw_np, self.p['sigma'], gain=self.p['gain'])
+                    raw_vst = raw_vst - bias
+
+                    ################# VST变换 #################
+                    lower = VST(0, self.p['sigma'], gain=self.p['gain'])
+                    upper = VST(self.p['scale'], self.p['sigma'], gain=self.p['gain'])
+                    nsr = 1 / (upper - lower)
+                    raw_vst = (raw_vst - lower) / (upper - lower)
+
+                    ################# 去噪 #################
+                    raw_dn = self.denoise(raw_vst, patch_size, nsr)
+                        
+                    ################# VST逆变换 #################
+                    raw_dn = raw_dn * (upper - lower) + lower
+                    self.denoised_data = inverse_VST(raw_dn, self.p['sigma'], gain=self.p['gain']) / self.p['scale']
+                    self.denoised_npy = rggb2bayer(self.denoised_data)
+
+                reg, self.mask_data = SimpleNLF(rggb2bayer(lr_raw), self.denoised_npy, k=13, 
                                                 setting={'mode':'collab', 'thr_mode':'score3', 'ransac': self.p['ransac']})
                 self.p['gain'] = reg[0] * self.p['scale']
                 self.p['sigma'] = np.sqrt(max(reg[1], 0)) * self.p['scale']
@@ -302,7 +317,6 @@ 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:
@@ -335,38 +349,11 @@ class YOND_Backend:
             raw_vst = (raw_vst - lower) / (upper - lower)
             
             ################# 准备去噪 #################
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            self.yond.net = self.yond.net.to(self.device)
-            raw_vst = torch.from_numpy(raw_vst).float().to(self.device).permute(2,0,1)[None,]
-            if 'guided' in self.yond.arch:
-                t = torch.tensor(nsr*self.p['sigsnr'], dtype=raw_vst.dtype, device=self.device)
-
-                # Denoise & pad
-                target_size = patch_size  # 设置目标块大小，可以根据需要调整 GPU:1024
-                overlap_ratio = 1/8  # 设置重叠率，可以根据需要调整
-                
-                # 使用改进的 big_image_split 函数
-                raw_inp, metadata = big_image_split(raw_vst, target_size, overlap_ratio)
-
-                raw_dn = torch.zeros_like(raw_inp[:,:4])
-                with torch.no_grad():
-                    if self.p['ddim_mode']:
-                        for i in range(raw_inp.shape[0]):  # 处理所有切块
-                            print(f'Patch: {i+1}/{len(raw_dn)}')
-                            raw_dn[i] = ddim(raw_inp[i][None,].clip(None, 2), self.yond.net, t, epoch=self.p['epoch'], 
-                                             sigma_t=self.p['sigma_t'], eta=self.p['eta_t'], sigma_corr=1.00)
-                    else:
-                        for i in range(raw_inp.shape[0]):  # 处理所有切块
-                            input_tensor = raw_inp[i][None,].clip(None, 2)
-                            raw_dn[i] = self.yond.net(input_tensor, t).clamp(0,None)
-
-                # 使用改进的 big_image_merge 函数
-                raw_dn = big_image_merge(raw_dn, metadata, blend_mode='avg')
+            raw_dn = self.denoise(raw_vst, patch_size, nsr)
             
             ################# VST逆变换 #################
-            raw_vst = raw_dn[0].permute(1,2,0).detach().cpu().numpy()
-            raw_vst = raw_vst * (upper - lower) + lower
-            self.denoised_data = inverse_VST(raw_vst, self.p['sigma'], gain=self.p['gain']) / self.p['scale']
+            raw_dn = raw_dn * (upper - lower) + lower
+            self.denoised_data = inverse_VST(raw_dn, self.p['sigma'], gain=self.p['gain']) / self.p['scale']
             
             self.denoised_npy = rggb2bayer(self.denoised_data)
             # 保存结果