Upload app.py

app.py

@@ -3,19 +3,24 @@ import numpy as np
 import torch
 from PIL import Image
 import gradio as gr
+from gradio_imageslider import ImageSlider
+import spaces  # 必须放在最前面，确保 ZeroGPU 初始化
+
+# 延迟 CUDA 初始化
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+weight_dtype = torch.float32
+
+# 加载模型组件
 from DAI.pipeline_all import DAIPipeline
 from DAI.controlnetvae import ControlNetVAEModel
 from DAI.decoder import CustomAutoencoderKL
 from diffusers import AutoencoderKL, UNet2DConditionModel
 from transformers import CLIPTextModel, AutoTokenizer
 
-# Initialize device and model paths
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-weight_dtype = torch.float32
 pretrained_model_name_or_path = "sjtu-deepvision/dereflection-any-image-v0"
 pretrained_model_name_or_path2 = "stabilityai/stable-diffusion-2-1"
 
-# Load the model components
+# 加载模型
 controlnet = ControlNetVAEModel.from_pretrained(pretrained_model_name_or_path, subfolder="controlnet", torch_dtype=weight_dtype).to(device)
 unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet", torch_dtype=weight_dtype).to(device)
 vae_2 = CustomAutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae_2", torch_dtype=weight_dtype).to(device)
@@ -23,7 +28,7 @@ vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path2, subfolder="v
 text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path2, subfolder="text_encoder").to(device)
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path2, subfolder="tokenizer", use_fast=False)
 
-# Create the pipeline
+# 创建推理管道
 pipe = DAIPipeline(
     vae=vae,
     text_encoder=text_encoder,
@@ -36,12 +41,13 @@ pipe = DAIPipeline(
     t_start=0,
 ).to(device)
 
-# Function to process the image
+# 使用 spaces.GPU 包装推理函数
+@spaces.GPU
 def process_image(input_image):
-    # Convert Gradio input to PIL Image
+    # 将 Gradio 输入转换为 PIL 图像
     input_image = Image.fromarray(input_image)
 
-    # Process the image
+    # 处理图像
     pipe_out = pipe(
         image=input_image,
         prompt="remove glass reflection",
@@ -49,16 +55,17 @@ def process_image(input_image):
         processing_resolution=None,
     )
 
-    # Convert the output to an image
+    # 将输出转换为图像
     processed_frame = (pipe_out.prediction.clip(-1, 1) + 1) / 2
     processed_frame = (processed_frame[0] * 255).astype(np.uint8)
     processed_frame = Image.fromarray(processed_frame)
 
-    return processed_frame
+    # 返回输入图像和处理后的图像
+    return input_image, processed_frame
 
-# Gradio interface
+# 创建 Gradio 界面
 def create_gradio_interface():
-    # Example images
+    # 示例图像
     example_images = [
         os.path.join("files", "image", f"{i}.png") for i in range(1, 9)
     ]
@@ -70,27 +77,33 @@ def create_gradio_interface():
                 input_image = gr.Image(label="Input Image", type="numpy")
                 submit_btn = gr.Button("Remove Reflection", variant="primary")
             with gr.Column():
-                output_image = gr.Image(label="Processed Image")
+                # 使用 ImageSlider 显示前后对比
+                output_slider = ImageSlider(
+                    label="Before & After",
+                    show_download_button=True,
+                    show_share_button=True,
+                )
 
-        # Add examples
+        # 添加示例
         gr.Examples(
             examples=example_images,
             inputs=input_image,
-            outputs=output_image,
+            outputs=output_slider,
             fn=process_image,
-            cache_examples=False,  # Cache results for faster loading
+            cache_examples=False,  # 缓存结果以加快加载速度
             label="Example Images",
         )
 
+        # 绑定按钮点击事件
         submit_btn.click(
             fn=process_image,
             inputs=input_image,
-            outputs=output_image,
+            outputs=output_slider,
         )
 
     return demo
 
-# Main function to launch the Gradio app
+# 主函数
 def main():
     demo = create_gradio_interface()
     demo.launch(server_name="0.0.0.0", server_port=7860)