Update app.py

app.py

@@ -1,22 +1,26 @@
-from transformers import pipeline
 import gradio as gr
+from PIL import Image
+from marigold_depth_estimation import MarigoldPipeline, UNet2DConditionModel, AutoencoderKL, DDIMScheduler
 
-# Initialize the pipeline with the Marigold model hosted on Hugging Face
-model = pipeline("image-to-image", model="prs-eth/marigold-depth-v1-0")
+# Instantiate the model components and the pipeline
+unet_model = UNet2DConditionModel()
+vae_model = AutoencoderKL()
+scheduler = DDIMScheduler()
+pipeline = MarigoldPipeline(unet=unet_model, vae=vae_model, scheduler=scheduler)
 
-def predict_depth(image):
-    # Generate a depth map from the input image using the model
-    output = model(image)
-    return output['output_image']  # Ensure this key matches the output of your model
+def predict_depth(input_image):
+    # Process the image and predict the depth map
+    output = pipeline(input_image)
+    return output.depth_image
 
-# Set up the Gradio interface
-interface = gr.Interface(
+iface = gr.Interface(
     fn=predict_depth,
-    inputs=gr.inputs.Image(shape=(512, 512), label="Upload Image"),
-    outputs=gr.outputs.Image(label="Depth Map"),
-    title="Marigold Depth Map Estimation",
-    description="Upload an image and the model will estimate and display its depth map."
+    inputs=gr.inputs.Image(type="pil", label="Upload an Image"),
+    outputs=gr.outputs.Image(type="pil", label="Depth Map"),
+    title="Depth Map Generation",
+    description="Upload an image to generate its depth map using the Marigold Depth Estimation Model.",
+    examples=["sample1.jpg", "sample2.jpg"]  # Optional: include example images in your repository
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()