Update app.py

app.py

@@ -1,98 +1,51 @@
-import gradio as gr
-import torch
-from transformers import AutoModel, AutoImageProcessor
-from PIL import Image
-
-# --- 1. 配置模型和设备 ---
-
-# 您的模型在 Hub 上的 ID
-repo_id = "suncongcong/AST_DeRain"
-
-# 自动选择使用 GPU 还是 CPU (Hugging Face Spaces 会提供免费的 CPU 或 T4 GPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"正在使用的设备: {device}")
-
-# --- 2. 加载模型和图像处理器 ---
-
-# 加载图像处理器
-processor = AutoImageProcessor.from_pretrained(repo_id)
-
-# 加载模型
-# trust_remote_code=True 是必须的，因为它需要执行您上传的 modeling_ast.py 文件
-model = AutoModel.from_pretrained(
-    repo_id,
-    trust_remote_code=True
-).to(device)
-
-print("模型加载成功，准备就绪！")
-
-
-# --- 3. 定义核心处理函数 ---
-
-def derain_image(input_image: Image.Image):
-    """
-    接收一个 PIL Image 对象，返回一个处理后的 PIL Image 对象。
-    """
-    # 确保输入是 RGB 格式
-    image = input_image.convert("RGB")
-
-    # 使用处理器将图片转换为模型所需的 Tensor 格式
-    inputs = processor(images=image, return_tensors="pt").to(device)
-
-    # 在不计算梯度的模式下进行推理，以节省显存和提高速度
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    # outputs[0] 就是我们得到的恢复后的图像 Tensor
-    restored_tensor = outputs[0] if isinstance(outputs, tuple) else outputs
-
-    # 将像素值限制在 0-1 之间
-    restored_tensor = torch.clamp(restored_tensor, 0, 1)
-
-    # 将 Tensor 转换回 PIL Image 以便显示
-    restored_image = processor.post_process_image_to_image(restored_tensor.cpu())[0]
-
-    return restored_image
-
-
-# --- 4. 创建并启动 Gradio 界面 ---
-
-# 使用 gr.Blocks() 可以更灵活地设计界面布局
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🖼️ AST 图像去雨模型在线演示
-        上传一张带雨的图片，模型将会自动去除雨水痕迹。
-        模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)
-        """
-    )
-
-    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,  # 按钮点击时调用的函数
-        inputs=input_img,  # 函数的输入来自哪个组件
-        outputs=output_img  # 函数的输出显示在哪个组件
-    )
-
-    # 添加一些示例图片，让用户可以快速体验
-    gr.Examples(
-        examples=[
-            ["https://i.imgur.com/a4y39hV.jpg"],
-            ["https://i.imgur.com/KxYE1v4.jpg"]
-        ],
-        inputs=input_img,
-        outputs=output_img,
-        fn=derain_image,
-        cache_examples=True  # 缓存示例结果，加快加载速度
-    )
-
-# 启动应用
-demo.launch()
+import gradio as gr
+import torch
+from transformers import CLIPImageProcessor
+from modeling_ast import ASTForRestoration
+from PIL import Image
+import requests
+from io import BytesIO
+from torchvision.transforms.functional import to_pil_image
+
+# --- 1. 配置模型和设备 ---
+repo_id = "suncongcong/AST_DeRain" 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+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)
+print("✅ 模型加载成功，准备就绪！")
+
+# --- 3. 定义核心处理函数 ---
+def derain_image(input_image: Image.Image):
+    if input_image is None: return None
+    image = input_image.convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    restored_tensor = outputs[0] if isinstance(outputs, tuple) else outputs
+    restored_tensor = torch.clamp(restored_tensor, 0, 1)
+    restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+    return restored_image
+
+# --- 4. 创建并启动 Gradio 界面 ---
+print("正在创建 Gradio Interface...")
+demo = gr.Interface(
+    fn=derain_image,
+    inputs=gr.Image(type="pil", label="输入带雨图片 (Input Rainy Image)"),
+    outputs=gr.Image(type="pil", label="输出清晰图片 (Output Deraided Image)"),
+    title="AST 图像去雨模型在线演示",
+    description="上传一张带雨的图片，模型将会自动去除雨水痕迹。模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)"
+)
+
+# --- 最终修正：添加 server_name 参数以适应容器环境 ---
+print("正在启动 Demo...")
+demo.launch(server_name="0.0.0.0")