File size: 1,854 Bytes
33baf83
 
072bd04
 
33baf83
072bd04
97677c5
072bd04
97677c5
 
072bd04
97677c5
 
 
 
 
 
072bd04
97677c5
 
072bd04
 
 
97677c5
 
072bd04
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
97677c5
 
 
 
 
 
41e6120
97677c5
072bd04
97677c5
072bd04
 
97677c5
 
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
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 minimal settings
        pipe = ShapEPipeline.from_pretrained(
            "openai/shap-e",
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True
        ).to("cpu")
        
        # Generate with minimal parameters
        outputs = pipe(
            prompt,
            num_inference_steps=16,
            frame_size=24,
            guidance_scale=7.5
        )
        
        # Get the mesh data
        mesh = outputs.meshes[0]  # Get the first mesh
        
        # Convert to vertices and faces format
        verts = mesh.verts.detach().cpu().numpy()
        faces = mesh.faces.detach().cpu().numpy()
        
        # Create trimesh object
        mesh_obj = trimesh.Trimesh(vertices=verts, faces=faces)
        
        # Export in desired format
        if output_path.endswith('.obj'):
            mesh_obj.export(output_path)
        elif output_path.endswith('.glb'):
            mesh_obj.export(output_path)
        elif output_path.endswith('.stl'):
            mesh_obj.export(output_path)
        
        print(f"Successfully exported 3D model to: {output_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 ring"
    try:
        # You can change the extension to .glb or .stl as needed
        generate_3d_model(prompt, "assistant_3d.obj")
    except Exception as e:
        print(f"Generation failed: {e}")