Update app.py

app.py

@@ -11,31 +11,62 @@ from torchvision.transforms.functional import to_pil_image, to_tensor
 from tqdm import tqdm
 
 # --- 1. 配置 ---
-repo_id = "suncongcong/AST_DeRain" 
+# 使用您提供的准确的模型仓库ID
+MODEL_IDS = {
+    "去雨痕 (Derain)": "Suncongcong/AST_DeRain",
+    "去雨滴 (Deraindrop)": "Suncongcong/AST_DeRainDrop",
+    "去雾 (Dehaze)": "Suncongcong/AST_Dehazing"
+}
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 PATCH_SIZE = 256
 OVERLAP = 64
 
 print(f"正在使用的设备: {device}")
 
-# --- 2. 加载模型和处理器 ---
-print(f"正在从 '{repo_id}' 加载模型和处理器...")
-processor = CLIPImageProcessor.from_pretrained(repo_id)
-processor.size = {"height": 256, "width": 256}
-processor.crop_size = {"height": 256, "width": 256}
-print(f"图像处理器尺寸已强制设置为: {processor.size}")
-model = ASTForRestoration.from_pretrained(
-    repo_id, 
-    trust_remote_code=True
-).to(device).eval()
-print("✅ 模型加载成功，准备就绪！")
-
-
-# --- 3. 定义“裁切-推理-合并”的核心处理函数 ---
-def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm=True)):
+# --- 2. 加载所有模型和处理器 ---
+MODELS = {}
+PROCESSOR = None
+
+print("正在加载所有模型和处理器...")
+# 使用 try-except 来增加鲁棒性
+try:
+    for task_name, repo_id in MODEL_IDS.items():
+        print(f"正在加载模型: {task_name} ({repo_id})")
+        if PROCESSOR is None:
+            PROCESSOR = CLIPImageProcessor.from_pretrained(repo_id)
+            print("✅ 处理器加载成功。")
+        
+        model = ASTForRestoration.from_pretrained(
+            repo_id, 
+            trust_remote_code=True
+        ).to(device).eval()
+        MODELS[task_name] = model
+        print(f"✅ 模型 '{task_name}' 加载成功。")
+except Exception as e:
+    print(f"加载模型时出错: {e}")
+    # 创建一个占位符函数，以便在模型加载失败时 Gradio 仍能启动并显示错误
+    def load_error_func(*args, **kwargs):
+        raise gr.Error(f"模型加载失败! 错误: {e}")
+    MODELS = {task: load_error_func for task in MODEL_IDS.keys()}
+
+
+print("所有模型加载完毕，准备就绪！")
+
+
+# --- 3. 定义统一的、可选择模型的处理函数 ---
+def process_image(input_image: Image.Image, task_name: str, progress=gr.Progress(track_tqdm=True)):
     if input_image is None:
         return None
-        
+    
+    # 根据传入的任务名称，选择对应的模型
+    model = MODELS[task_name]
+    print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}")
+    
+    # 检查模型是否加载成功
+    if not isinstance(model, torch.nn.Module):
+        model() # 这会触发上面定义的错误函数
+
     img = input_image.convert("RGB")
     img_tensor = to_tensor(img).unsqueeze(0).to(device)
     b, c, h, w = img_tensor.shape
@@ -49,7 +80,7 @@ def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm
     w_steps = len(range(0, w, stride))
     total_patches = h_steps * w_steps
     
-    pbar = tqdm(total=total_patches, desc="正在处理图像块...")
+    pbar = tqdm(total=total_patches, desc=f"正在执行 {task_name}...")
 
     for y in range(0, h, stride):
         for x in range(0, w, stride):
@@ -61,7 +92,7 @@ def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm
             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), 'replicate') # <-- 最终修正
+                patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'replicate')
             else:
                 patch_padded = patch_in
             
@@ -81,26 +112,35 @@ def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm
     pbar.close()
 
     restored_tensor = output_canvas / weight_map
-    
     restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
     
     return restored_image
 
-# --- 4. 创建并启动 Gradio 界面 ---
+# --- 4. 创建并启动带选项卡的 Gradio 界面 ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        #AST 图像去雨模型在线演示
-        上传任意尺寸的带雨图片，模型将会分块处理并拼接成完整的高清输出。
-        模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)
+        # 🖼️ 多功能图像复原工具 (AST 模型)
+        请选择一个任务，然后上传对应的图片进行处理。
         """
     )
-    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)
     
+    with gr.Tabs():
+        # 根据 MODEL_IDS 字典自动创建选项卡
+        for task_name in MODEL_IDS.keys():
+            with gr.TabItem(task_name, id=task_name):
+                with gr.Row():
+                    input_img = gr.Image(type="pil", label=f"输入图片 (Input for {task_name})")
+                    output_img = gr.Image(type="pil", label="输出图片 (Output)")
+                
+                task_id_box = gr.Textbox(task_name, visible=False)
+                
+                submit_btn = gr.Button("开始处理 (Process)", variant="primary")
+                
+                submit_btn.click(
+                    fn=process_image, 
+                    inputs=[input_img, task_id_box], 
+                    outputs=output_img
+                )
+
 demo.launch(server_name="0.0.0.0")