Update app.py

app.py

@@ -1,29 +1,137 @@
 import gradio as gr
-from diffusers import MarigoldDepthPipeline, DDIMScheduler
+import os
+import torch
+
 import torch
 from PIL import Image
+from diffusers import (
+    AutoencoderKL,
+)
 
-CHECKPOINT = "developy/ApDepth"
+from transformers import CLIPTextModel, CLIPTokenizer
+from apdepth import MarigoldPipeline
+from apdepth.modules.unet_2d_condition import UNet2DConditionModel
 
-device = "cpu"
-dtype = torch.float32
+def load_example(example_image):
+    # 返回选中的图片
+    return example_image
 
-pipe = MarigoldDepthPipeline.from_pretrained(CHECKPOINT)
-pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-pipe = pipe.to(device=device, dtype=dtype)
 
-def predict(image: Image.Image):
-    out = pipe(image)
-    depth_vis = pipe.image_processor.visualize_depth(out.prediction)[0]
-    return depth_vis
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_repo_id = "developy/ApDepth"  # Replace to the model you would like to use
+
+torch_dtype = torch.float32
+
+vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae", torch_dtype=torch_dtype, allow_pickle=False)
+unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder="unet", torch_dtype=torch_dtype, allow_pickle=False)
+text_encoder = CLIPTextModel.from_pretrained(model_repo_id, subfolder="text_encoder", torch_dtype=torch_dtype)
+tokenizer = CLIPTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer", torch_dtype=torch_dtype)
+pipe = DepthMasterPipeline(vae=vae, unet=unet, text_encoder=text_encoder, tokenizer=tokenizer)
+
+
+try:
+    pipe.enable_xformers_memory_efficient_attention()
+except ImportError:
+    pass  # run without xformers
+
+pipe = pipe.to(device)
+
+
+# @spaces.GPU #[uncomment to use ZeroGPU]
+def infer(
+    input_image,
+    progress=gr.Progress(track_tqdm=True),
+):
+
+    pipe_out = pipe(
+            input_image,
+            processing_res=768,
+            match_input_res=True,
+            batch_size=1,
+            color_map="Spectral",
+            show_progress_bar=True,
+            resample_method="bilinear",
+        )
+
+    # depth_pred: np.ndarray = pipe_out.depth_np
+    depth_colored: Image.Image = pipe_out.depth_colored
+
+
+    return depth_colored
+
+
+# 默认图像路径
+example_images = [
+    "example/example_0.jpg",
+    "example/example_1.jpg",
+    "example/example_2.jpg",
+    "example/example_3.jpg",
+    "example/example_4.jpg",
+    "example/example_5.jpg",
+    "example/example_6.jpg"
+]
+
+# css = """
+# #col-container {
+#     margin: 0 auto;
+#     max-width: 640px;
+# }
+# #example-gallery {
+#     height: 80px; /* 设置缩略图高度 */
+#     width: auto;  /* 保持宽高比 */
+#     margin: 0 auto;  /* 图片间距 */
+#     cursor: pointer; /* 鼠标指针变为手型 */
+# }
+# """
+
+css = """
+#img-display-container {
+    max-height: 100vh;
+}
+#img-display-input {
+    max-height: 80vh;
+}
+#img-display-output {
+    max-height: 80vh;
+}
+#download {
+    height: 62px;
+}
+"""
+
+title = "# ApDepth"
+description = """**Official demo for ApDepth**.
+Please refer to our [website](https://haruko386.github.io/research/) for more details."""
+
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    gr.Markdown(" ### Depth Estimation with ApDepth.")
+    # with gr.Column(elem_id="col-container"):
+    #     gr.Markdown(" # Depth Estimation")
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type="pil", elem_id="img-display-input")
+        with gr.Column():
+            # depth_img_slider = ImageSlider(label="Depth Map with Slider View", elem_id="img-display-output", position=0.5)
+            depth_map = gr.Image(label="Depth Map with Slider View", type="pil", interactive=False, elem_id="depth-map")
+
+    # 计算按钮
+    compute_button = gr.Button(value="Compute Depth")
+
+    # 设置计算按钮的回调
+    compute_button.click(
+        fn=infer,  # 回调函数
+        inputs=[input_image],  # 输入
+        outputs=[depth_map]  # 输出
+    )
+
+    example_files = os.listdir('example')
+    example_files.sort()
+    example_files = [os.path.join('example', filename) for filename in example_files]
+    examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[depth_map], fn=infer)
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Depth Map"),
-    title="ApDepth Demo",
-    description="Monocular Depth Estimation based on Marigold"
-)
 
-if __name__ == "__main__":
-    demo.launch()
+# 启动 Gradio 应用
+demo.queue().launch(share=True)