import gradio as gr
from PIL import Image
from generate import generate_model_from_text, generate_model_from_image

def from_text(prompt):
    output = generate_model_from_text(prompt)
    return output

def from_image(img: Image.Image):
    output = generate_model_from_image(img)
    return output

with gr.Blocks() as demo:
    gr.Markdown("# 🧊 3D Model Generator\nGenerate downloadable 3D models from text or image")

    with gr.Tab("Text to 3D"):
        txt = gr.Textbox(label="Enter a prompt")
        txt_btn = gr.Button("Generate")
        txt_out = gr.File(label="Download 3D Model")
        txt_btn.click(from_text, inputs=txt, outputs=txt_out)

    with gr.Tab("Image to 3D"):
        img = gr.Image(type="pil")
        img_btn = gr.Button("Generate")
        img_out = gr.File(label="Download 3D Model")
        img_btn.click(from_image, inputs=img, outputs=img_out)

demo.launch()