File size: 2,099 Bytes
31f6f71 | 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 | import argparse
import os
os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
import numpy as np
from PIL import Image
from trellis.pipelines import (
DVDImageToVoxelPipeline,
TrellisImageTo3DPipeline,
export_cubified_voxels,
run_image_stage2_from_dvd_voxels,
)
from trellis.utils import postprocessing_utils
def parse_args():
parser = argparse.ArgumentParser(description="DVD image-to-voxel generation followed by TRELLIS stage 2.")
parser.add_argument("--image", required=True, help="Input image path.")
parser.add_argument("--dvd-config", default="ckpts/dvd_img.json", help="DVD model config JSON.")
parser.add_argument("--dvd-checkpoint", default="ckpts/dvd_img.safetensors", help="DVD safetensors checkpoint.")
parser.add_argument("--output-dir", default="example_results", help="Directory for generated assets.")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device", default="cuda")
return parser.parse_args()
def main():
args = parse_args()
os.makedirs(args.output_dir, exist_ok=True)
image = Image.open(args.image)
name = os.path.splitext(os.path.basename(args.image))[0]
dvd = DVDImageToVoxelPipeline.from_files(
args.dvd_config,
args.dvd_checkpoint,
device=args.device,
)
voxels = dvd.sample_voxels(image, seed=args.seed)
voxel_coords_path = os.path.join(args.output_dir, f"voxel64_{name}_dis.npy")
np.save(voxel_coords_path, voxels.coords_without_batch.numpy())
export_cubified_voxels(voxels, os.path.join(args.output_dir, f"{name}_dvd_voxels.glb"))
trellis = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
trellis.to(args.device)
outputs = run_image_stage2_from_dvd_voxels(trellis, image, voxels, seed=args.seed)
glb = postprocessing_utils.to_glb(outputs["gaussian"][0], outputs["mesh"][0])
glb.export(os.path.join(args.output_dir, f"{name}.glb"))
print(f"Saved DVD coords: {voxel_coords_path}")
if __name__ == "__main__":
main()
|