import torch
import numpy as np
import gradio as gr
import trimesh
from skimage import measure
from PIL import Image
from diffusers import StableDiffusionPipeline
Load the Stable Diffusion model (make sure the model supports GPU if available)
def generate_3d_model(prompt: str):
"""
Generates a 3D model (OBJ file) from a text prompt.
This function:
1. Generates an image using Stable Diffusion.
2. Converts the image to grayscale to simulate a depth map.
3. Applies an inversion to simulate depth.
4. Uses marching cubes to create a 3D mesh.
5. Exports the mesh to an OBJ file.
"""
# Generate image from prompt
image = pipe(prompt).images[0]
# Convert image to grayscale for depth simulation
image_gray = image.convert("L")
depth_array = 255 - np.array(image_gray)
# Normalize depth for marching cubes
depth_norm = depth_array / 255.0
# Generate mesh using marching cubes
try:
verts, faces, _, _ = measure.marching_cubes(depth_norm, level=0.5)
except Exception as e:
return f"Error generating mesh: {e}"
mesh = trimesh.Trimesh(vertices=verts, faces=faces)
# Save mesh as OBJ file
output_path = "generated_model.obj"
mesh.export(output_path)
return output_path
Create a Gradio interface
interface = gr.Interface(
fn=generate_3d_model,
inputs=gr.Textbox(lines=2, placeholder="Enter a prompt (e.g., futuristic spaceship)", label="Prompt"),
outputs=gr.File(label="Download 3D Model (OBJ)"),
title="AI-Powered 3D Model Generator",
description="Generate a 3D model from a text prompt using Stable Diffusion and a diffusion-based depth estimation."
)
