File size: 1,422 Bytes
b119edc a778497 b119edc a778497 b119edc a778497 b119edc a778497 b119edc a778497 b119edc a778497 b119edc a778497 b119edc a778497 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
import base64
import os
from PIL import Image
from io import BytesIO
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import postprocessing_utils
class EndpointHandler:
def __init__(self, model_dir):
# Load the pipeline from the local directory
self.pipeline = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
self.pipeline.cuda()
def __call__(self, data):
"""
Args:
data (:obj:`dict`):
- "inputs": The base64 encoded image or URL.
- "params": Dictionary of generation parameters (optional).
"""
inputs = data.pop("inputs", data)
params = data.pop("parameters", {})
# Decode image
image = Image.open(BytesIO(base64.b64decode(inputs)))
# Run Inference
# Note: You can adjust 'steps' or 'cfg' via params
outputs = self.pipeline(
image,
num_samples=1,
return_flags=["mesh"],
**params
)
# Process mesh to GLB
mesh = outputs['mesh'][0]
glb_io = BytesIO()
mesh.export(glb_io, file_type='glb')
glb_io.seek(0)
# Encode GLB to base64 for the response
return {
"mesh_base64": base64.b64encode(glb_io.getvalue()).decode("utf-8"),
"format": "glb"
}
|