File size: 2,276 Bytes
33baf83 072bd04 33baf83 072bd04 97677c5 072bd04 97677c5 377ce24 97677c5 377ce24 97677c5 377ce24 97677c5 377ce24 97677c5 377ce24 072bd04 377ce24 072bd04 377ce24 4a3b38f 377ce24 4a3b38f 377ce24 4a3b38f 072bd04 97677c5 41e6120 97677c5 083ed22 97677c5 072bd04 97677c5 377ce24 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
from diffusers import ShapEPipeline
import trimesh
import numpy as np
def generate_3d_model(prompt, output_path="assistant_3d.obj"):
"""
Generate a 3D model using ShapE and export it in a Blender-compatible format
"""
try:
# Load pipeline with memory constraints suitable for CPU-only usage
pipe = ShapEPipeline.from_pretrained(
"openai/shap-e",
torch_dtype=torch.float32,
low_cpu_mem_usage=True # Minimize CPU memory usage
).to("cpu")
# Generate with reduced inference steps and guidance scale for CPU
outputs = pipe(
prompt=prompt,
num_inference_steps=8, # Reduced steps for CPU speed
guidance_scale=5.0, # Lower guidance scale to save memory
)
# Extract vertices and faces for the 3D mesh generation
vertices = outputs["vertices"][0].detach().cpu().numpy()
faces = outputs["faces"][0].detach().cpu().numpy()
# Create the trimesh object with error handling
mesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
# Export model in desired format with robust handling for compatibility
try:
if output_path.endswith('.obj'):
mesh_obj.export(output_path, include_normals=True)
else:
mesh_obj.export(output_path)
print(f"Successfully exported 3D model to: {output_path}")
except Exception as export_error:
print(f"Error during export: {export_error}")
# Use a fallback format in case of export failure
backup_path = output_path.rsplit('.', 1)[0] + '.ply'
mesh_obj.export(backup_path)
print(f"Exported backup model to: {backup_path}")
return output_path
except Exception as e:
print(f"Error during generation: {e}")
print(f"Error type: {type(e)}")
print(f"Full error details: {str(e)}")
raise
if __name__ == "__main__":
prompt = "a simple 3D ring, perfect circle, clean geometry"
try:
generate_3d_model(prompt, "assistant_3d.obj")
except Exception as e:
print(f"Generation failed: {e}")
|