import torch
import numpy as np
from diffusers import DiffusionPipeline
import streamlit as st

# Load the ShapE pipeline on CPU
pipeline = DiffusionPipeline.from_pretrained(
    "openai/shap-e",
    torch_dtype=torch.float32,
    trust_remote_code=True,
    custom_pipeline="openai/shap-e",  # Assuming it works with custom_pipeline param
).to("cpu")

# Define the function to generate and save a 3D model
def generate_3d_model(prompt, output_path="/tmp/output.ply"):
    # Run the pipeline with a text prompt
    result = pipeline(prompt, None)
    
    # Try to save the result as a 3D model
    try:
        pipeline.save_ply(result, output_path)
        print(f"Model saved to {output_path}")
        return output_path
    except Exception as e:
        print(f"Failed to save model: {e}")
        return None

# Streamlit interface
st.title("3D Model Generator")
prompt = st.text_input("Enter a prompt to generate a 3D model:", "a cat statue")

if st.button("Generate Model"):
    with st.spinner("Generating model..."):
        model_path = generate_3d_model(prompt)
        if model_path:
            st.success("Model generated successfully!")
            # Display download link
            with open(model_path, "rb") as file:
                st.download_button(
                    label="Download 3D Model",
                    data=file,
                    file_name="generated_model.ply",
                    mime="application/octet-stream"
                )
        else:
            st.error("Model generation failed.")