Enhance generate3d function to initialize nvdiffrast context for mesh export, improving rendering capabilities. Updated export call to include rendering context and adjusted parameters for UVs and texture image, ensuring better integration with the model.
db2fd1d
| import os | |
| import torch | |
| from mesh import Mesh | |
| # Ensure the output directory exists | |
| def ensure_dir(path): | |
| os.makedirs(path, exist_ok=True) | |
| def generate3d(model, rgb_image, xyz_image, device="cpu"): | |
| output_dir = "outputs" | |
| ensure_dir(output_dir) | |
| prompt_id = "mesh_output" | |
| base_path = os.path.join(output_dir, prompt_id) | |
| obj_path = base_path + ".obj" | |
| glb_path = base_path + ".glb" | |
| # CRM export expects a data dictionary | |
| data = {"rgb": rgb_image, "xyz": xyz_image} | |
| # Get rendering context required by xatlas and nvdiffrast | |
| try: | |
| import nvdiffrast.torch as dr | |
| ctx = dr.RasterizeGLContext() | |
| except Exception as e: | |
| raise RuntimeError("Failed to initialize nvdiffrast context: " + str(e)) | |
| # Export mesh with UVs and texture image | |
| model.export_mesh_wt_uv(ctx, data, base_path, 0, device, 512) | |
| # Convert .obj to .glb | |
| mesh = Mesh.load(obj_path, device=torch.device("cpu")) | |
| mesh.write(glb_path) | |
| return glb_path | |