import gradio as gr
from PIL import Image
from diffusers import StableDiffusionLDM3DPipeline

# Load the model. Do this once to avoid reloading on every request.
pipe = StableDiffusionLDM3DPipeline.from_pretrained("Intel/ldm3d-pano")


def generate_images(prompt, guidance_scale=7.0, num_inference_steps=50):
    output = pipe(
        prompt,
        width=1024,
        height=512,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
    )

    rgb_image, depth_image = output.rgb, output.depth

    # Convert to PIL Images for Gradio compatibility
    rgb_image = Image.fromarray(rgb_image[0])
    depth_image = Image.fromarray(depth_image[0])

    return rgb_image, depth_image


iface = gr.Interface(
    fn=generate_images,
    inputs=[
        "text",
        gr.Slider(0, 20, value=7.0, label="Guidance Scale"),
        gr.Slider(0, 100, value=50, label="Inference Steps")
    ],
    outputs=["image", "image"],
    title="ldm3d-pano Image Generator"
)

iface.launch()