Upload app.py

app.py

@@ -1,109 +1,18 @@
-from __future__ import annotations
-
-import functools
 import os
-import tempfile
-import logging
-
 import gradio as gr
-import numpy as np
-
 from PIL import Image
 from gradio_imageslider import ImageSlider
-from tqdm import tqdm
-
-from pathlib import Path
-import gradio
-from gradio.utils import get_cache_folder
-from DAI.pipeline_all import DAIPipeline
-
-from diffusers import (
-    AutoencoderKL,
-    UNet2DConditionModel,
-)
-
-from transformers import CLIPTextModel, AutoTokenizer
-
-from DAI.controlnetvae import ControlNetVAEModel
-
-from DAI.decoder import CustomAutoencoderKL
-
-import torch 
-
-
-class Examples(gradio.helpers.Examples):
-    def __init__(self, *args, directory_name=None, **kwargs):
-        super().__init__(*args, **kwargs, _initiated_directly=False)
-        if directory_name is not None:
-            self.cached_folder = get_cache_folder() / directory_name
-            self.cached_file = Path(self.cached_folder) / "log.csv"
-        self.create()
-
-
-default_seed = 2024
-default_batch_size = 1
-
-
-def process_image_check(path_input):
-    logging.info(f"Input image path: {path_input}")
-    if path_input is None:
-        raise gr.Error(
-            "Missing image in the first pane: upload a file or use one from the gallery below."
-        )
-
-def resize_image(input_image, resolution):
-    if not isinstance(input_image, Image.Image):
-        raise ValueError("input_image should be a PIL Image object")
 
-    input_image_np = np.asarray(input_image)
-    H, W, C = input_image_np.shape
-    H = float(H)
-    W = float(W)
-    
-    k = float(resolution) / min(H, W)
-    
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    
-    img = input_image.resize((W, H), Image.Resampling.LANCZOS)
-    
-    return img
 
-def process_image(
-    pipe,
-    vae_2,
-    path_input,
-):
+def process_image_direct(path_input):
     try:
-        name_base, name_ext = os.path.splitext(os.path.basename(path_input))
-        logging.info(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")
         input_image = Image.open(path_input)
-
-        pipe_out = pipe(
-            image=input_image,
-            prompt="remove glass reflection",
-            vae_2=vae_2,
-            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)
-        yield [input_image, path_out_png]
+        return [input_image, path_input]
     except Exception as e:
-        logging.error(f"Error processing image: {e}")
-        yield [None, None]
+        return [None, None]
 
 
-def run_demo_server(pipe, vae_2):
-    process_pipe_image = functools.partial(process_image, pipe, vae_2)
-
+def run_demo_server():
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(
@@ -127,7 +36,7 @@ def run_demo_server(pipe, vae_2):
             """
             # Dereflection Any Image
             <p align="center">
-            Upload an image to remove reflections.
+            Upload an image to display it directly.
             </p>
             """
         )
@@ -142,7 +51,7 @@ def run_demo_server(pipe, vae_2):
                         )
                         with gr.Row():
                             image_submit_btn = gr.Button(
-                                value="Remove Reflection", variant="primary"
+                                value="Display Image", variant="primary"
                             )
                             image_reset_btn = gr.Button(value="Reset")
                     with gr.Column():
@@ -154,8 +63,7 @@ def run_demo_server(pipe, vae_2):
                             elem_classes="slider",
                         )
 
-                Examples(
-                    fn=process_pipe_image,
+                gr.Examples(
                     examples=sorted([
                         os.path.join("files", "image", name)
                         for name in os.listdir(os.path.join("files", "image"))
@@ -163,23 +71,13 @@ def run_demo_server(pipe, vae_2):
                     ]),
                     inputs=[image_input],
                     outputs=[image_output_slider],
-                    cache_examples=True,
-                    directory_name="examples_image",
+                    cache_examples=False,
                 )
 
         image_submit_btn.click(
-            fn=process_image_check,
+            fn=process_image_direct,
             inputs=image_input,
-            outputs=None,
-            preprocess=False,
-            queue=False,
-        ).success(
-            fn=process_pipe_image,
-            inputs=[
-                image_input,
-            ],
             outputs=[image_output_slider],
-            concurrency_limit=1,
         )
 
         image_reset_btn.click(
@@ -192,49 +90,10 @@ def run_demo_server(pipe, vae_2):
                 image_input,
                 image_output_slider,
             ],
-            queue=False,
-        )
-
-        demo.queue(api_open=True).launch(share=False)
-
-
-def 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
-
-    controlnet = ControlNetVAEModel.from_pretrained(pretrained_model_name_or_path, subfolder="controlnet", torch_dtype=torch.float32)
-    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet", torch_dtype=torch.float32)
-    vae_2 = CustomAutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae_2", torch_dtype=torch.float32)
-
-    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,
         )
 
-    run_demo_server(pipe, vae_2)
+        demo.launch(share=False)
 
 
 if __name__ == "__main__":
-    main()
+    run_demo_server()