import gradio as gr
from PIL import Image
from marigold_depth_estimation import MarigoldPipeline, UNet2DConditionModel, AutoencoderKL, DDIMScheduler

# 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(input_image):
    # Process the image and predict the depth map
    output = pipeline(input_image)
    return output.depth_image

iface = gr.Interface(
    fn=predict_depth,
    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__":
    iface.launch()