import shutil
import gradio as gr
from gradio_client import Client

# Define a function to generate and return the .glb file path
def generate_3d_model(prompt):
    # Initialize Gradio client for Shap-E API
    client = Client("hysts/Shap-E")
    result = client.predict(
        prompt=prompt,
        seed=0,
        guidance_scale=15,
        num_inference_steps=64,
        api_name="/text-to-3d"
    )

    # Define the path to save the .glb file
    glb_path = "generated_model.glb"

    # Copy the result file locally if it's a temp file
    if result.startswith("/tmp"):
        shutil.copy(result, glb_path)
    else:
        # Otherwise, download the file if it's a URL
        import requests
        response = requests.get(result)
        with open(glb_path, "wb") as f:
            f.write(response.content)

    return glb_path

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_3d_model,
    inputs="text",
    outputs=gr.Model3D(label="Generated 3D Model"),
    title="3D Model Generator",
    description="Enter a prompt to generate a 3D model."
)

# Launch the Gradio app
interface.launch()