Refactor background removal process in app.py to utilize rembg library, enhancing performance and simplifying the code. Update device handling to allow dynamic selection between CPU and CUDA, improving compatibility across different hardware configurations. Modify output format from OBJ to GLB for better integration with Gradio display.
9d0b3b4
| import os | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from mesh import Mesh | |
| from pipelines.pipeline_text_to_3d import TextTo3D | |
| # === Load Model (assumes this is done once at startup, not per request) === | |
| model = TextTo3D.from_pretrained("./checkpoints/zeroscope_v1_5") | |
| model.to(torch.device("cpu")) | |
| model.eval() | |
| def generate3d(prompt: str, guidance_scale: float = 15.0, steps: int = 50) -> str: | |
| # === Set up paths === | |
| output_dir = "outputs" | |
| os.makedirs(output_dir, exist_ok=True) | |
| base_name = prompt.replace(" ", "_").lower() | |
| mesh_path_base = os.path.join(output_dir, base_name) | |
| # === Generate 3D Mesh === | |
| mesh = model(prompt, guidance_scale=guidance_scale, steps=steps) | |
| obj_path = mesh_path_base + ".obj" | |
| mesh.export_mesh_wt_uv(obj_path) | |
| # === Convert to GLB with textures === | |
| mesh_loaded = Mesh.load(obj_path, device=torch.device("cpu")) | |
| glb_path = mesh_path_base + ".glb" | |
| mesh_loaded.write(glb_path) | |
| # === Return GLB path for Gradio display === | |
| return glb_path | |
| if __name__ == "__main__": | |
| # Example run | |
| prompt = "a modern wooden chair" | |
| output_glb = generate3d(prompt) | |
| print(f"Generated GLB: {output_glb}") | |