Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,543 Bytes
8fff686 c948747 8fff686 c948747 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
# Project EmbodiedGen
#
# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
from embodied_gen.utils.monkey_patch.trellis import monkey_path_trellis
monkey_path_trellis()
import random
import torch
from PIL import Image
from embodied_gen.data.utils import trellis_preprocess
from embodied_gen.models.sam3d import Sam3dInference
from embodied_gen.utils.trender import pack_state, unpack_state
from thirdparty.TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline
__all__ = [
"image3d_model_infer",
]
def image3d_model_infer(
pipe: TrellisImageTo3DPipeline | Sam3dInference,
seg_image: Image.Image,
seed: int = None,
**kwargs: dict,
) -> dict[str, any]:
"""Execute 3D generation using Trellis or SAM3D pipeline on input image."""
if isinstance(pipe, TrellisImageTo3DPipeline):
pipe.cuda()
seg_image = trellis_preprocess(seg_image)
outputs = pipe.run(
seg_image,
preprocess_image=False,
seed=(random.randint(0, 100000) if seed is None else seed),
# Optional parameters
# sparse_structure_sampler_params={
# "steps": 12,
# "cfg_strength": 7.5,
# },
# slat_sampler_params={
# "steps": 12,
# "cfg_strength": 3,
# },
**kwargs,
)
pipe.cpu()
elif isinstance(pipe, Sam3dInference):
outputs = pipe.run(
seg_image,
seed=(random.randint(0, 100000) if seed is None else seed),
# stage1_inference_steps=25,
# stage2_inference_steps=25,
**kwargs,
)
state = pack_state(outputs["gaussian"][0], outputs["mesh"][0])
# Align GS3D from SAM3D with TRELLIS format.
outputs["gaussian"][0], _ = unpack_state(state, device="cuda")
else:
raise ValueError(f"Unsupported pipeline type: {type(pipe)}")
torch.cuda.empty_cache()
return outputs
|