app.py

import spaces  # 必须放在最前面
import os
import numpy as np
import torch
from PIL import Image
import gradio as gr

# 延迟 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

pretrained_model_name_or_path = "sjtu-deepvision/dereflection-any-image-v0"
pretrained_model_name_or_path2 = "stabilityai/stable-diffusion-2-1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载模型
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)
vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path2, subfolder="vae").to(device)
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)

# 创建推理管道
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,
).to(device)

@spaces.GPU
def process_image(input_image, resolution_option):
    # 将 Gradio 输入转换为 PIL 图像
    input_image = Image.fromarray(input_image)

    # 根据用户选择的分辨率选项设置 processing_resolution
    processing_resolution = 2560 if resolution_option == "2560" else None

    # 处理图像
    pipe_out = pipe(
        image=input_image,
        prompt="remove glass reflection",
        vae_2=vae_2,
        processing_resolution=processing_resolution,
    )

    # 将输出转换为图像
    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

# 创建 Gradio 界面
def create_gradio_interface():
    # 示例图像
    example_images = [
        os.path.join("files", "image", f"{i}.png") for i in range(1, 9)
    ]
    title = "# Dereflection Any Image"
    description = """Official demo for **Dereflection Any Image**.
    Please refer to our [paper](), [project page](https://abuuu122.github.io/DAI.github.io/), and [github](https://github.com/Abuuu122/Dereflection-Any-Image) for more details."""

    with gr.Blocks() as demo:
        gr.Markdown(title)
        gr.Markdown(description)

        # 创建两个 Tab
        with gr.Tabs():
            with gr.Tab("Image"):
                with gr.Row():
                    with gr.Column():
                        input_image_default = gr.Image(label="Input Image", type="numpy")
                        submit_btn_default = gr.Button("Remove Reflection", variant="primary")
                    with gr.Column():
                        output_image_default = gr.Image(label="Processed Image")

                # 添加示例
                gr.Examples(
                    examples=example_images,
                    inputs=input_image_default,
                    outputs=output_image_default,
                    fn=lambda x: process_image(x, "None"),
                    cache_examples=False,  # 缓存结果以加快加载速度
                    label="Example Images",
                )

                # 绑定按钮点击事件
                submit_btn_default.click(
                    fn=lambda x: process_image(x, "None"),
                    inputs=input_image_default,
                    outputs=output_image_default,
                )

            with gr.Tab("2K (2560)"):
                with gr.Row():
                    with gr.Column():
                        input_image_high = gr.Image(label="Input Image", type="numpy")
                        submit_btn_high = gr.Button("Remove Reflection", variant="primary")
                    with gr.Column():
                        output_image_high = gr.Image(label="Processed Image")

                # 添加示例
                gr.Examples(
                    examples=example_images,
                    inputs=input_image_high,
                    outputs=output_image_high,
                    fn=lambda x: process_image(x, "2560"),
                    cache_examples=False,  # 缓存结果以加快加载速度
                    label="Example Images",
                )

                # 绑定按钮点击事件
                submit_btn_high.click(
                    fn=lambda x: process_image(x, "2560"),
                    inputs=input_image_high,
                    outputs=output_image_high,
                )

    return demo

# 主函数
def main():
    demo = create_gradio_interface()
    demo.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    main()