Run TripoSG 2D-to-3D on beam.cloud

beam_triposg.py

#
# Run TripoSG 2D-to-3D on beam.cloud
#
# Luke Van In (x.com/lukevanin)
#
# This script is used to deploy the TripoSG 2D-to-3D model on beam.cloud.
#
# The model is deployed as a cloud endpoint, which can be called directly
# using the beam.cloud API.
#
# The model is deployed using the beam.cloud CLI, which can be installed
# using the following command:
#
# Log in to beam.cloud and generate a token (MY_TOKEN).
#
# Clone the TripoSG repo:
# git clone https://github.com/VAST-AI/TripoSG.git
# cd TripoSG
#
# Then, install the beam-client:
# brew install beam-client
# pip install --upgrade beam
# beam configure default --token MY_TOKEN
#
# Create volumes for images and models:
# beam volume create images
# beam volume create models
#
# The model can then be deployed using the following command:
#
# beam deploy beam_triposg.py:generate_mesh --name tripo-sg
#
# This will return a URL for the endpoint.The model can then be called 
# using the following command:
# YOUR_IMAGE_BASE64 is the base64 encoded image,
# YOUR_IMAGE_NAME is the name of the image file (stored in the images volume).
# YOUR_MODEL_NAME is the name of the model file (stored in the models volume).
#
# curl -X POST <endpoint_url> \
#     -H "Authorization: Bearer YOUR_BEAM_TOKEN" \
#     -H "Content-Type: application/json" \
#     -d '{"image_base64": "YOUR_IMAGE_BASE64", "image_name": "YOUR_IMAGE_NAME", "model_name": "YOUR_MODEL_NAME"}'
#
# This will return a URL which can be used to download the model:
#
# curl -L <endpoint_url> -o <model_name>.glb
#


import base64
import beam, asyncio
from beam import Output, endpoint
from beam import Image as BeamImage
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import Response

# ─────────────── 1. Build the image ───────────────
triposg_image = (
    beam.Image(
        base_image="nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04",
        python_version="python3.10",
    )
    .add_python_packages(
        [
            # CUDA‑enabled PyTorch + TripoSG deps
            "accelerate",
            "torch",
            "torchvision",
            "diffusers==0.30.3",
            "transformers==4.37.2",
            "einops",
            "huggingface_hub",
            "opencv-python",
            "scikit-image",
            "trimesh",
            "omegaconf",
            "numpy==1.22.3",
            "jaxtyping",
            "typeguard",
            "pillow",
            "fastapi[standard]==0.115.4",
            "https://huggingface.co/spaces/VAST-AI/TripoSG/resolve/main/"
            "diso-0.1.4-cp310-cp310-linux_x86_64.whl",
        ]
    )
    .add_commands(
        [
            # Manually install GL dependencies 
            # clang and build-essential may not be required
            "apt-get update && apt-get install -y libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender1",
        ]
    )
)




# ─────────────── 2. GPU function ───────────────
@endpoint(
    image=triposg_image,
    gpu="A100-40",
    timeout=600,
    volumes=[
        beam.Volume(name="images", mount_path="/volumes/images"),
        beam.Volume(name="models", mount_path="/volumes/models"),
    ]
)
def generate_mesh(image_base64: str, image_name: str, model_name: str):

    import sys, io, os, inspect, shutil
    from PIL import Image
    import numpy as np
    import torch
    import trimesh
    from huggingface_hub import snapshot_download

    # sys.path.append("/TripoSG")          # repo cloned at build time
    from triposg.pipelines.pipeline_triposg import TripoSGPipeline
    import triposg.schedulers.scheduling_rectified_flow as rf

    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
    sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/scripts")

    from briarmbg import BriaRMBG
    from inference_triposg import run_triposg


    device = "cuda"
    dtype = torch.float16
    seed = 42
    num_inference_steps = 50
    guidance_scale = 7.0


    print("Decoding image...", len(image_base64))
    image_bytes = base64.b64decode(image_base64)

    # Save image to file
    with open(f"/volumes/images/{image_name}", "wb") as f:
        f.write(image_bytes)

    # download pretrained weights
    triposg_weights_dir = "pretrained_weights/TripoSG"
    rmbg_weights_dir = "pretrained_weights/RMBG-1.4"

    # Only download if the directory doesn't exist
    if not os.path.exists(triposg_weights_dir):
        print(f"Downloading TripoSG weights to {triposg_weights_dir}...")
        snapshot_download(repo_id="VAST-AI/TripoSG", local_dir=triposg_weights_dir)
    else:
        print(f"TripoSG weights already downloaded to {triposg_weights_dir}")

    if not os.path.exists(rmbg_weights_dir):
        print(f"Downloading RMBG weights to {rmbg_weights_dir}...")
        snapshot_download(repo_id="briaai/RMBG-1.4", local_dir=rmbg_weights_dir)
    else:
        print(f"RMBG weights already downloaded to {rmbg_weights_dir}")

    # init rmbg model for background removal
    rmbg_net = BriaRMBG.from_pretrained(rmbg_weights_dir).to(device)
    rmbg_net.eval() 

    # init tripoSG pipeline
    pipe: TripoSGPipeline = TripoSGPipeline.from_pretrained(triposg_weights_dir).to(device, dtype)

    # run inference
    print("Running inference...")
    mesh = run_triposg(
        pipe,
        image_input=f"/volumes/images/{image_name}",
        rmbg_net=rmbg_net,
        seed=seed,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    )

    model_path = os.path.join("/volumes/models", model_name)
    print(f"Exporting mesh to {model_path}...")
    mesh.export(model_path)

    print(f"Saving output to {model_path}...")
    output = Output(path=model_path)
    output.save()
    output_url = output.public_url()

    print(f"Output saved to {output_url}")
    return output_url