File size: 1,539 Bytes
33baf83 072bd04 f051324 33baf83 f051324 97677c5 f051324 072bd04 97677c5 f051324 41e6120 f051324 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
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.")
|