Initial TripoSG Space app

.gitignore

@@ -0,0 +1,5 @@
+checkpoints/
+triposg/
+tmp/
+__pycache__/
+*.glb

README.md

@@ -1,12 +1,51 @@
 ---
-title: Trip W Oblaka
+title: TripoSG Image-to-3D API
 emoji: 😻
 colorFrom: pink
 colorTo: blue
 sdk: gradio
 sdk_version: 6.5.0
 app_file: app.py
+python_version: 3.10
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TripoSG Image-to-3D API
+
+This Space wraps the official TripoSG pipeline and exposes a `/predict` API endpoint for programmatic generation of GLB meshes.
+
+## API usage
+
+Python:
+
+```python
+from gradio_client import Client
+
+client = Client("your-username/your-space")
+result = client.predict(
+    image_path="input.png",
+    seed=0,
+    num_inference_steps=50,
+    guidance_scale=7.5,
+    simplify=True,
+    target_face_num=100000,
+    api_name="/predict",
+)
+print(result)
+```
+
+Raw HTTP (example):
+
+```bash
+curl -X POST \
+  -H "Content-Type: application/json" \
+  -d '{"data": ["data:image/png;base64,......", 0, 50, 7.5, true, 100000]}' \
+  https://your-username-your-space.hf.space/api/predict
+```
+
+The response contains the generated GLB file path and URL.
+
+## Notes
+
+- The Space will clone `VAST-AI-Research/TripoSG` at runtime and download weights from `VAST-AI/TripoSG` and `briaai/RMBG-1.4`.
+- `requirements.txt` targets the default Hugging Face Spaces GPU runtime (Linux). For local runs, adjust Torch/CUDA and the `diso` wheel as needed.

app.py

@@ -0,0 +1,225 @@
+import os
+import sys
+import uuid
+import shutil
+
+import gradio as gr
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+import trimesh
+
+try:
+    import spaces
+
+    gpu = spaces.GPU
+except Exception:
+
+    def gpu(*_args, **_kwargs):
+        def _wrap(fn):
+            return fn
+
+        return _wrap
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
+TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
+TRIPOSG_CODE_DIR = "./triposg"
+
+CHECKPOINT_DIR = "checkpoints"
+RMBG_PRETRAINED_MODEL = os.path.join(CHECKPOINT_DIR, "RMBG-1.4")
+TRIPOSG_PRETRAINED_MODEL = os.path.join(CHECKPOINT_DIR, "TripoSG")
+
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
+os.makedirs(TMP_DIR, exist_ok=True)
+
+
+if not os.path.exists(TRIPOSG_CODE_DIR):
+    os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}")
+
+sys.path.append(TRIPOSG_CODE_DIR)
+sys.path.append(os.path.join(TRIPOSG_CODE_DIR, "scripts"))
+
+
+from image_process import prepare_image
+from briarmbg import BriaRMBG
+from triposg.pipelines.pipeline_triposg import TripoSGPipeline
+from utils import simplify_mesh
+
+
+snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
+rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE)
+rmbg_net.eval()
+
+snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL)
+triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(
+    DEVICE, DTYPE
+)
+
+
+def _session_dir(req: gr.Request | None) -> str:
+    if req is None:
+        return TMP_DIR
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(save_dir, exist_ok=True)
+    return save_dir
+
+
+def _unique_glb_path(save_dir: str) -> str:
+    return os.path.join(save_dir, f"triposg_{uuid.uuid4().hex}.glb")
+
+
+def _run_triposg(
+    image_path: str,
+    seed: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    simplify: bool,
+    target_face_num: int,
+    req: gr.Request | None = None,
+):
+    if not image_path:
+        raise gr.Error("Upload an image first.")
+
+    image_seg = prepare_image(
+        image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net
+    )
+
+    generator = torch.Generator(device=triposg_pipe.device).manual_seed(seed)
+    outputs = triposg_pipe(
+        image=image_seg,
+        generator=generator,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+    ).samples[0]
+
+    mesh = trimesh.Trimesh(outputs[0].astype(np.float32), np.ascontiguousarray(outputs[1]))
+
+    if simplify:
+        mesh = simplify_mesh(mesh, target_face_num)
+
+    save_dir = _session_dir(req)
+    mesh_path = _unique_glb_path(save_dir)
+    mesh.export(mesh_path)
+
+    return image_seg, mesh_path
+
+
+@gpu(duration=180)
+@torch.no_grad()
+def generate_mesh(
+    image_path: str,
+    seed: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    simplify: bool,
+    target_face_num: int,
+    req: gr.Request | None = None,
+):
+    image_seg, mesh_path = _run_triposg(
+        image_path,
+        seed,
+        num_inference_steps,
+        guidance_scale,
+        simplify,
+        target_face_num,
+        req,
+    )
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return image_seg, mesh_path
+
+
+@gpu(duration=180)
+@torch.no_grad()
+def api_generate(
+    image_path: str,
+    seed: int,
+    num_inference_steps: int,
+    guidance_scale: float,
+    simplify: bool,
+    target_face_num: int,
+    req: gr.Request | None = None,
+):
+    _, mesh_path = _run_triposg(
+        image_path,
+        seed,
+        num_inference_steps,
+        guidance_scale,
+        simplify,
+        target_face_num,
+        req,
+    )
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return mesh_path
+
+
+def _cleanup_session(req: gr.Request):
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+TITLE = "TripoSG Image-to-3D API"
+DESCRIPTION = (
+    "Upload a single-object image to generate a 3D mesh (GLB). "
+    "This demo exposes a /predict API endpoint."
+)
+
+
+with gr.Blocks(title=TITLE) as demo:
+    gr.Markdown(f"# {TITLE}\n\n{DESCRIPTION}")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Input Image", type="filepath")
+            seg_output = gr.Image(
+                label="Segmentation Preview", type="pil", format="png"
+            )
+
+            with gr.Accordion("Generation Settings", open=True):
+                seed = gr.Slider(
+                    label="Seed", minimum=0, maximum=2**31 - 1, step=1, value=0
+                )
+                steps = gr.Slider(
+                    label="Inference Steps", minimum=8, maximum=50, step=1, value=50
+                )
+                guidance = gr.Slider(
+                    label="CFG Scale", minimum=0.0, maximum=20.0, step=0.1, value=7.5
+                )
+                simplify = gr.Checkbox(label="Simplify Mesh", value=True)
+                face_count = gr.Slider(
+                    label="Target Face Count",
+                    minimum=10000,
+                    maximum=1000000,
+                    step=1000,
+                    value=100000,
+                )
+
+            generate_btn = gr.Button("Generate 3D", variant="primary")
+
+        with gr.Column():
+            model_output = gr.Model3D(label="Generated GLB", interactive=False)
+            file_output = gr.File(label="Download GLB", interactive=False)
+
+    generate_btn.click(
+        generate_mesh,
+        inputs=[image_input, seed, steps, guidance, simplify, face_count],
+        outputs=[seg_output, model_output],
+    ).then(lambda path: path, inputs=model_output, outputs=file_output)
+
+    api_btn = gr.Button(visible=False)
+    api_btn.click(
+        api_generate,
+        inputs=[image_input, seed, steps, guidance, simplify, face_count],
+        outputs=[file_output],
+        api_name="/predict",
+    )
+
+    demo.unload(_cleanup_session)
+
+
+demo.launch()

requirements.txt

@@ -0,0 +1,23 @@
+diffusers==0.32.2
+trimesh
+pillow
+spandrel==0.4.0
+plyfile==1.1
+xformers
+pymcubes==0.1.4
+shapely
+mkl==2022.0.2
+nvdiffrast
+cvcuda_cu12==0.6.0.16
+triton==3.1.0
+imageio==2.36.0
+numpy==1.26.4
+scipy==1.13.1
+tqdm==4.67.1
+opencv-python
+open3d==0.18.0
+pymeshlab
+ninja==1.11.1.3
+matplotlib
+
+diso @ https://github.com/Chumbyte/DiSO/releases/download/v0.1.4/diso-0.1.4-cp310-cp310-linux_x86_64.whl

utils.py

@@ -0,0 +1,37 @@
+import numpy as np
+import open3d as o3d
+import pymeshlab as pml
+
+
+def simplify_mesh(mesh, target_face_num: int = 100000):
+    if mesh.faces.shape[0] <= target_face_num:
+        return mesh
+
+    vertices = mesh.vertices
+    faces = mesh.faces
+
+    ms = pml.MeshSet()
+    ms.add_mesh(pml.Mesh(vertices, faces))
+    ms.meshing_decimation_quadric_edge_collapse(
+        targetfacenum=int(target_face_num), preserveboundary=True
+    )
+
+    new_mesh = ms.current_mesh()
+    new_vertices = new_mesh.vertex_matrix()
+    new_faces = new_mesh.face_matrix()
+
+    o3d_mesh = o3d.geometry.TriangleMesh(
+        o3d.utility.Vector3dVector(new_vertices),
+        o3d.utility.Vector3iVector(new_faces),
+    )
+    o3d_mesh = o3d_mesh.remove_duplicated_vertices()
+    o3d_mesh = o3d_mesh.remove_degenerate_triangles()
+    o3d_mesh = o3d_mesh.remove_non_manifold_edges()
+    o3d_mesh = o3d_mesh.remove_unreferenced_vertices()
+
+    return mesh.__class__(
+        vertices=np.asarray(o3d_mesh.vertices),
+        faces=np.asarray(o3d_mesh.triangles),
+        vertex_normals=np.asarray(o3d_mesh.vertex_normals),
+        process=False,
+    )