app.py

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}")