Upload app.py

app.py

@@ -1,100 +1,110 @@
-import gradio as gr
-from PIL import Image
-from DAI.pipeline_all import DAIPipeline
+import spaces  # 必须放在最前面
 import os
-import tempfile
 import numpy as np
+import torch
+from PIL import Image
+import gradio as gr
+from gradio_imageslider import ImageSlider
 
-from diffusers import (
-    AutoencoderKL,
-    UNet2DConditionModel,
-)
-
-from transformers import CLIPTextModel, AutoTokenizer
+# 延迟 CUDA 初始化
+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
 
-def process_image(pipe, vae_2, image):
-    # Save the input image to a temporary file
-    temp_input_path = tempfile.mktemp(suffix=".png")
-    image.save(temp_input_path)
+pretrained_model_name_or_path = "sjtu-deepvision/dereflection-any-image-v0"
+pretrained_model_name_or_path2 = "stabilityai/stable-diffusion-2-1"
+
+# 加载模型
+controlnet = ControlNetVAEModel.from_pretrained(pretrained_model_name_or_path, subfolder="controlnet", torch_dtype=weight_dtype)
+unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet", torch_dtype=weight_dtype)
+vae_2 = CustomAutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae_2", torch_dtype=weight_dtype)
+vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path2, subfolder="vae")
+text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path2, subfolder="text_encoder")
+tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path2, subfolder="tokenizer", use_fast=False)
+
+# 创建推理管道
+pipe = DAIPipeline(
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    unet=unet,
+    controlnet=controlnet,
+    safety_checker=None,
+    scheduler=None,
+    feature_extractor=None,
+    t_start=0,
+)
 
-    name_base, name_ext = os.path.splitext(os.path.basename(temp_input_path))
-    print(f"Processing image {name_base}{name_ext}")
 
-    path_output_dir = tempfile.mkdtemp()
-    path_out_png = os.path.join(path_output_dir, f"{name_base}_delight.png")
-    resolution = None
+def process_image(input_image):
+    # 将 Gradio 输入转换为 PIL 图像
+    input_image = Image.fromarray(input_image)
 
+    # 处理图像
     pipe_out = pipe(
-        image=image,
+        image=input_image,
         prompt="remove glass reflection",
         vae_2=vae_2,
-        processing_resolution=resolution,
+        processing_resolution=None,
     )
 
+    # 将输出转换为图像
     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)
-    processed_frame.save(path_out_png)
 
-    return processed_frame
+    # 返回输入图像和处理后的图像
+    return input_image, processed_frame
+
+# 创建 Gradio 界面
+def create_gradio_interface():
+    # 示例图像
+    example_images = [
+        os.path.join("files", "image", f"{i}.png") for i in range(1, 9)
+    ]
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Dereflection Any Image")
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="numpy")
+                submit_btn = gr.Button("Remove Reflection", variant="primary")
+            with gr.Column():
+                # 使用 ImageSlider 显示前后对比
+                output_slider = ImageSlider(
+                    label="Before & After",
+                    show_download_button=True,
+                    show_share_button=True,
+                )
+
+        # 添加示例
+        gr.Examples(
+            examples=example_images,
+            inputs=input_image,
+            outputs=output_slider,
+            fn=process_image,
+            cache_examples=False,  # 缓存结果以加快加载速度
+            label="Example Images",
+        )
+
+        # 绑定按钮点击事件
+        submit_btn.click(
+            fn=process_image,
+            inputs=input_image,
+            outputs=output_slider,
+        )
+
+    return demo
+
+# 主函数
+def main():
+    demo = create_gradio_interface()
+    demo.launch(server_name="0.0.0.0", server_port=7860)
 
 if __name__ == "__main__":
-    pretrained_model_name_or_path = "JichenHu/dereflection-any-image-v0"
-    pretrained_model_name_or_path2 = "stabilityai/stable-diffusion-2-1"
-    revision = None
-    variant = None
-
-    # Load the model
-    controlnet = ControlNetVAEModel.from_pretrained(pretrained_model_name_or_path, subfolder="controlnet")
-    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
-    vae_2 = CustomAutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae_2")
-
-    vae = AutoencoderKL.from_pretrained(
-        pretrained_model_name_or_path2, subfolder="vae", revision=revision, variant=variant
-    )
-
-    text_encoder = CLIPTextModel.from_pretrained(
-        pretrained_model_name_or_path2, subfolder="text_encoder", revision=revision, variant=variant
-    )
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_name_or_path2,
-        subfolder="tokenizer",
-        revision=revision,
-        use_fast=False,
-    )
-    pipe = DAIPipeline(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        controlnet=controlnet,
-        safety_checker=None,
-        scheduler=None,
-        feature_extractor=None,
-        t_start=0,
-    )
-
-    # Cache example images in memory
-    example_images_dir = "files/image"
-    example_images = []
-    for i in range(1, 9):
-        image_path = os.path.join(example_images_dir, f"{i}.png")
-        if os.path.exists(image_path):
-            example_images.append([Image.open(image_path)])
-
-    # Create a Gradio interface
-    interface = gr.Interface(
-        fn=lambda image: process_image(pipe, vae_2, image),
-        inputs=gr.Image(type="pil"),
-        outputs=gr.Image(type="pil"),
-        title="Dereflection Any Image",
-        description="Upload an image to remove glass reflections.",
-        examples=example_images,
-    )
-
-    interface.launch()
-
+    main()