Refactor 3D mesh generation in inference.py to streamline the process. Introduced ensure_dir function for output directory management, updated generate3d function to accept model and image inputs, and improved mesh export and conversion to GLB format. This enhances code clarity and maintains compatibility with Gradio display.

inference.py

@@ -1,39 +1,25 @@
 import os
 import torch
-import numpy as np
-from PIL import Image
 from mesh import Mesh
-from pipelines.pipeline_text_to_3d import TextTo3D
 
+# Ensure the output directory exists
+def ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
 
-# === 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 ===
+def generate3d(model, rgb_image, xyz_image, device="cpu"):
     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)
+    ensure_dir(output_dir)
 
-    # === 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)
+    prompt_id = "mesh_output"
+    base_path = os.path.join(output_dir, prompt_id)
+    obj_path = base_path + ".obj"
+    glb_path = base_path + ".glb"
 
-    # === Return GLB path for Gradio display ===
-    return glb_path
+    # Export mesh using CRM model (generates .obj and textures)
+    model.export_mesh_wt_uv(obj_path, rgb=rgb_image, xyz=xyz_image, device=device)
 
+    # Load and convert to GLB
+    mesh = Mesh.load(obj_path, device=torch.device("cpu"))
+    mesh.write(glb_path)
 
-if __name__ == "__main__":
-    # Example run
-    prompt = "a modern wooden chair"
-    output_glb = generate3d(prompt)
-    print(f"Generated GLB: {output_glb}")
+    return glb_path