Upload app.py

app.py

@@ -1,18 +1,90 @@
 import gradio as gr
 from PIL import Image
+from DAI.pipeline_all import DAIPipeline
+import os
+import tempfile
+import numpy as np
 
-def process_image(image):
-    # Simply return the input image without any processing
-    return image
-
-# Create a Gradio interface
-interface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Simple Image Echo App",
-    description="Upload an image and get the same image as output."
+from diffusers import (
+    AutoencoderKL,
+    UNet2DConditionModel,
 )
 
+from transformers import CLIPTextModel, AutoTokenizer
+
+from DAI.controlnetvae import ControlNetVAEModel
+
+from DAI.decoder import CustomAutoencoderKL
+
+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)
+
+    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
+
+    pipe_out = pipe(
+        image=image,
+        prompt="remove glass reflection",
+        vae_2=vae_2,
+        processing_resolution=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)
+    processed_frame.save(path_out_png)
+
+    return processed_frame
+
 if __name__ == "__main__":
-    interface.launch()
+    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,
+    )
+
+    # 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="Image Dereflection App",
+        description="Upload an image to remove glass reflections."
+    )
+
+    interface.launch()