Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-import torch.nn.functional as F 
+import torch.nn.functional as F
 import numpy as np
 from transformers import CLIPImageProcessor
 from modeling_ast import ASTForRestoration
@@ -35,21 +35,20 @@ print("✅ 模型加载成功，准备就绪！")
 def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm=True)):
     if input_image is None:
         return None
-
+        
     img = input_image.convert("RGB")
     img_tensor = to_tensor(img).unsqueeze(0).to(device)
     b, c, h, w = img_tensor.shape
 
     output_canvas = torch.zeros_like(img_tensor).to(device)
     weight_map = torch.zeros_like(img_tensor).to(device)
-
+    
     stride = PATCH_SIZE - OVERLAP
-
-    # 计算需要裁切的块数
+    
     h_steps = len(range(0, h, stride))
     w_steps = len(range(0, w, stride))
     total_patches = h_steps * w_steps
-
+    
     pbar = tqdm(total=total_patches, desc="正在处理图像块...")
 
     for y in range(0, h, stride):
@@ -57,34 +56,34 @@ def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm
             y_end = min(y + PATCH_SIZE, h)
             x_end = min(x + PATCH_SIZE, w)
             patch_in = img_tensor[:, :, y:y_end, x:x_end]
-
+            
             ph, pw = patch_in.shape[2:]
             pad_h = PATCH_SIZE - ph
             pad_w = PATCH_SIZE - pw
             if pad_h > 0 or pad_w > 0:
-                patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'reflect')
+                patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'replicate') # <-- 最终修正
             else:
                 patch_padded = patch_in
-
+            
             with torch.no_grad():
                 outputs = model(patch_padded)
-
+            
             patch_out = outputs[0] if isinstance(outputs, tuple) else outputs
             patch_out = torch.clamp(patch_out, 0, 1)
 
             patch_out_unpadded = patch_out[:, :, :ph, :pw]
-
+            
             output_canvas[:, :, y:y_end, x:x_end] += patch_out_unpadded
             weight_map[:, :, y:y_end, x:x_end] += 1
-
+            
             pbar.update(1)
 
     pbar.close()
 
     restored_tensor = output_canvas / weight_map
-
+    
     restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
-
+    
     return restored_image
 
 # --- 4. 创建并启动 Gradio 界面 ---
@@ -99,9 +98,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         input_img = gr.Image(type="pil", label="输入带雨图片 (Input Rainy Image)")
         output_img = gr.Image(type="pil", label="输出清晰图片 (Output Deraided Image)")
-
+        
     submit_btn = gr.Button("开始去雨 (Start Deraining)", variant="primary")
-
+    
     submit_btn.click(fn=derain_image_Tiled, inputs=input_img, outputs=output_img)
-
+    
 demo.launch(server_name="0.0.0.0")