import os import json from easydict import EasyDict as edict from trellis import models, datasets, trainers import copy from torch.utils.data import DataLoader import torch from PIL import Image import numpy as np from IPython.display import display import torch import numpy as np import random import time from trellis.datasets.sparse_structure_latent import SparseStructureLatentVisMixin from trellis.pipelines import TrellisImageTo3DPipeline from trellis.utils import render_utils, postprocessing_utils import rembg import utils3d from trellis.renderers import OctreeRenderer from trellis.representations import Octree import imageio import os from trellis.modules import sparse as sp import argparse from tqdm import tqdm from subprocess import DEVNULL, call, TimeoutExpired import math import queue import threading def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) def _render(file_path, engine='CYCLES', cuda_idx=0, total_render_view_num=300, eval_view_num=10, BLENDER_EEVEE_NEXT_on_0_device=False, timeout_seconds=600, no_norm_scene=False, debug=False, blender_path=''): output_folder = os.path.join(os.path.dirname(file_path), 'render') os.makedirs(output_folder, exist_ok=True) if os.path.exists(os.path.join(output_folder, 'transforms.json')): return output_folder # Build camera {yaw, pitch, radius, fov} (shared with GS eval branch) yaws, pitchs, radius_f, fov_deg = _eval_view_cameras(total_render_view_num, eval_view_num) radius = [radius_f] * eval_view_num fov = [fov_deg / 180 * np.pi] * eval_view_num views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': fv} for y, p, r, fv in zip(yaws, pitchs, radius, fov)] args = [ blender_path, '-b', '-P', './dataset_toolkits/blender_script/render.py', '--', '--views', json.dumps(views), '--object', os.path.expanduser(file_path), '--resolution', '512', '--output_folder', output_folder, '--engine', engine, # ('BLENDER_EEVEE_NEXT', 'CYCLES') ] if no_norm_scene: args.append('--no_norm_scene') if file_path.endswith('.blend'): args.insert(1, file_path) if BLENDER_EEVEE_NEXT_on_0_device: args.insert(1, '--gpu-backend') args.insert(2, 'vulkan') try: if engine == 'CYCLES': env = os.environ.copy() env['CUDA_VISIBLE_DEVICES'] = str(cuda_idx) if debug: call(args, env=env, timeout=timeout_seconds) else: call(args, env=env, stdout=DEVNULL, timeout=timeout_seconds) else: if debug: call(args, timeout=timeout_seconds) else: call(args, stdout=DEVNULL, timeout=timeout_seconds) except TimeoutExpired: print(f"Render timed out for {output_folder} using {engine}") return output_folder def _eval_view_cameras(total_render_view_num=300, eval_view_num=10): """ Subsampled yaw/pitch trajectory used by both Blender `_render()` and GS eval views. Matches `torch.linspace` / sine pitch schedule, then `:: step` subsampling. Returns yaw/pitch lists (rad), scalar radius, and vertical FOV in degrees (for `render_utils`). """ yaws = torch.linspace(0, 2 * 3.1415, total_render_view_num) pitchs = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, total_render_view_num)) step = total_render_view_num // eval_view_num yaws = yaws[::step].tolist() pitchs = pitchs[::step].tolist() radius = 2.0 fov_deg = 40.0 return yaws, pitchs, radius, fov_deg def _render_gaussian_eval_views( gaussian_sample, output_folder, total_render_view_num=300, eval_view_num=10, resolution=512, bg_color=(0, 0, 0), verbose=False, ): """ Rasterize 3D Gaussians from decoded SLAT with the same cameras as mesh Blender rendering. Writes `000.png` ... under `output_folder` and `transforms.json` (metadata only). """ os.makedirs(output_folder, exist_ok=True) marker = os.path.join(output_folder, 'transforms.json') if os.path.exists(marker): return output_folder yaws, pitchs, r, fov_deg = _eval_view_cameras(total_render_view_num, eval_view_num) extrinsics, intrinsics = render_utils.yaw_pitch_r_fov_to_extrinsics_intrinsics( yaws, pitchs, r, fov_deg ) rets = render_utils.render_frames( gaussian_sample, extrinsics, intrinsics, options={'resolution': resolution, 'bg_color': bg_color}, verbose=verbose, ) colors = rets['color'] for i, img in enumerate(colors): Image.fromarray(img).save(os.path.join(output_folder, f'{i:03d}.png')) fov_rad = float(np.deg2rad(fov_deg)) frames = [] for i, (y, p) in enumerate(zip(yaws, pitchs)): frames.append({ 'file_path': f'{i:03d}.png', 'camera_angle_x': fov_rad, 'yaw': y, 'pitch': p, 'radius': r, }) with open(marker, 'w') as f: json.dump({'backend': 'gaussian_splat', 'frames': frames}, f, indent=2) return output_folder # process image def preprocess_image(input: Image.Image, no_crop=False) -> Image.Image: """ Preprocess the input image. """ # if has alpha channel, use it directly; otherwise, remove background has_alpha = False if input.mode == 'RGBA': alpha = np.array(input)[:, :, 3] if not np.all(alpha == 255): has_alpha = True if has_alpha: output = input else: input = input.convert('RGB') max_size = max(input.size) scale = min(1, 1024 / max_size) if scale < 1: input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS) rembg_session = rembg.new_session('u2net') output = rembg.remove(input, session=rembg_session) output_np = np.array(output) alpha = output_np[:, :, 3] if not no_crop: bbox = np.argwhere(alpha > 0.8 * 255) bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) size = int(size * 1.2) bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2 output = output.crop(bbox) # type: ignore output = output.resize((518, 518), Image.Resampling.LANCZOS) output = np.array(output).astype(np.float32) / 255 output = output[:, :, :3] * output[:, :, 3:4] output = Image.fromarray((output * 255).astype(np.uint8)) return output def render_image_list(x_0, num_frames=300, resolution=256): renderer = OctreeRenderer() renderer.rendering_options.resolution = resolution renderer.rendering_options.near = 0.8 renderer.rendering_options.far = 1.6 renderer.rendering_options.bg_color = (0, 0, 0) renderer.rendering_options.ssaa = 4 renderer.pipe.primitive = 'voxel' # Build camera yaws = torch.linspace(0, 2 * 3.1415, num_frames) pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) yaws = yaws.tolist() pitch = pitch.tolist() exts = [] ints = [] for yaw, pitch in zip(yaws, pitch): orig = torch.tensor([ np.sin(yaw) * np.cos(pitch), np.cos(yaw) * np.cos(pitch), np.sin(pitch), ]).float().cuda() * 2 fov = torch.deg2rad(torch.tensor(30)).cuda() extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov) exts.append(extrinsics) ints.append(intrinsics) images = [] # Build each representation x_0 = x_0.cuda() assert x_0.shape[0] == 1 i = 0 representation = Octree( depth=10, aabb=[-0.5, -0.5, -0.5, 1, 1, 1], device='cuda', primitive='voxel', sh_degree=0, primitive_config={'solid': True}, ) coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False) resolution = x_0.shape[-1] representation.position = coords.float() / resolution representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(resolution)), dtype=torch.uint8, device='cuda') for _, (ext, intr) in enumerate(zip(exts, ints)): res = renderer.render(representation, ext, intr, colors_overwrite=representation.position) images.append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) return images class MyDataset(SparseStructureLatentVisMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.normalization = None self.loads = [100 for _ in range(100)] def __len__(self): return 100 def collate_fn(self, batch): return batch class TrellisEdititngModel: def __init__(self, ss_latents_config_path, ss_latents_load_dir, ss_latents_load_ckpt, latents_config_path, latents_load_dir, latents_load_ckpt, load_ema_model_for_inference=False, trellis_pipeline_path='/mnt/zsn/ckpts/TRELLIS-image-large'): self.trellis_pipeline_path = trellis_pipeline_path self.ss_latents_config_path = ss_latents_config_path self.ss_latents_load_dir = ss_latents_load_dir self.ss_latents_load_ckpt = ss_latents_load_ckpt self.latents_config_path = latents_config_path self.latents_load_dir = latents_load_dir self.latents_load_ckpt = latents_load_ckpt # load ss_latents model config = json.load(open(self.ss_latents_config_path, 'r')) cfg = edict() cfg.update(config) cfg.data_dir = '/home/dataset_model/dataset/3DGS/objaverse_v1' cfg.output_dir = '../../work_dirs/Editing_Training/Debug' cfg.load_dir = self.ss_latents_load_dir cfg.load_ckpt = self.ss_latents_load_ckpt if load_ema_model_for_inference: cfg.trainer.args.load_ema_model_for_inference = True model_dict = { name: getattr(models, model.name)(**model.args).cuda() for name, model in cfg.models.items() } self.dataset = MyDataset() self.ss_latents_trainer = getattr(trainers, cfg.trainer.name)(model_dict, self.dataset, **cfg.trainer.args, output_dir=cfg.output_dir, load_dir=cfg.load_dir, step=cfg.load_ckpt) self.ss_latents_sampler = self.ss_latents_trainer.get_sampler() # load latents model config = json.load(open(self.latents_config_path, 'r')) cfg = edict() cfg.update(config) cfg.data_dir = '/home/dataset_model/dataset/3DGS/objaverse_v1' cfg.output_dir = '../../work_dirs/Editing_Training/Debug' cfg.load_dir = self.latents_load_dir cfg.load_ckpt = self.latents_load_ckpt if load_ema_model_for_inference: cfg.trainer.args.load_ema_model_for_inference = True model_dict = { name: getattr(models, model.name)(**model.args).cuda() for name, model in cfg.models.items() } self.latents_trainer = getattr(trainers, cfg.trainer.name)(model_dict, self.dataset, **cfg.trainer.args, output_dir=cfg.output_dir, load_dir=cfg.load_dir, step=cfg.load_ckpt) self.latents_sampler = self.latents_trainer.get_sampler() # ori pipeline self.ori_pipeline = TrellisImageTo3DPipeline.from_pretrained(self.trellis_pipeline_path) self.ori_pipeline.cuda() # other params self.total_steps = 25 self.gs_bg_color = (0, 0, 0) def editing_img_to_ss_latents(self, edited_img): pass def editing_img_to_latents(self, edited_img): pass def editing_inference( self, ori_ss_latents_path, ori_latents_path, edited_img_path, ori_img_path, vis_resolution=256, ori_latents_norm=False, edited_ss_latents_path=None, video_save_path=None, mesh_save_path=None, slat_save_path=None, render_gs_views=False, gs_render_total_view_num=300, gs_render_eval_view_num=10, gs_render_resolution=512, output_video=True, output_mesh=True, print_time=False, empty_structure_fallback='error', ): with torch.no_grad(): assert output_video or output_mesh or slat_save_path or render_gs_views # load img processed_ori_img = preprocess_image(Image.open(ori_img_path), no_crop=False) processed_ori_img_tensor = torch.from_numpy(np.array(processed_ori_img)).permute(2, 0, 1).unsqueeze(0).cuda() / 255.0 processed_img = preprocess_image(Image.open(edited_img_path), no_crop=False) processed_img_tensor = torch.from_numpy(np.array(processed_img)).permute(2, 0, 1).unsqueeze(0).cuda() / 255.0 # load ss_latents ori_ss_latent = torch.from_numpy(np.load(ori_ss_latents_path)['mean']).cuda()[None] ori_latents_feats = torch.from_numpy(np.load(ori_latents_path)['feats']).cuda() if ori_latents_norm: std = torch.tensor(self.ori_pipeline.slat_normalization['std'])[None].to(ori_latents_feats.device) mean = torch.tensor(self.ori_pipeline.slat_normalization['mean'])[None].to(ori_latents_feats.device) ori_latents_feats = (ori_latents_feats - mean) / std # print(torch.mean(ori_latents_feats, dim=0), torch.std(ori_latents_feats, dim=0)) ori_latents_coords = torch.from_numpy(np.load(ori_latents_path)['coords']).cuda() if ori_latents_coords.shape[1] == 3: ori_latents_coords = torch.cat([torch.zeros_like(ori_latents_coords[:, :1]), ori_latents_coords], dim=1) # preparing data data_tensor = dict() data_tensor['ori_cond_img'] = processed_ori_img_tensor data_tensor['edited_cond_img'] = processed_img_tensor data_tensor['ori_ss_latent'] = ori_ss_latent data_tensor['edited_ss_latent'] = ori_ss_latent args = self.ss_latents_trainer.get_inference_cond(**data_tensor) args['cond'] = torch.clone(args['edited_cond_img']) time_a = time.time() # inference ss_latents if edited_ss_latents_path is None: ss_latents_noise = torch.randn_like(ori_ss_latent) res = self.ss_latents_sampler.sample( self.ss_latents_trainer.models['denoiser'], noise=ss_latents_noise, **args, steps=self.total_steps, cfg_strength=0, rescale_t=3.0, start_step=0, end_step=self.total_steps, verbose=False, ).samples else: res = torch.from_numpy(np.load(edited_ss_latents_path)['mean']).cuda()[None] time_b = time.time() if print_time: print(f'------SS latents inference time: {time_b - time_a} seconds') # get coords voxel = self.latents_trainer.dataset.decode_latent(res) > 0 coords = torch.argwhere(voxel)[:, [0, 2, 3, 4]].int() if coords.shape[0] == 0: if empty_structure_fallback == 'original': print('WARNING: empty sparse structure predicted; falling back to original latent coords') coords = ori_latents_coords.int() else: raise RuntimeError('empty sparse structure predicted') noise = sp.SparseTensor( feats=torch.randn(coords.shape[0], self.latents_trainer.models['denoiser'].in_channels).cuda(), coords=coords, ) ori_latents = sp.SparseTensor( feats=ori_latents_feats, coords=ori_latents_coords.int(), ) # preparing data data_tensor['ori_ss_latent'] = ori_latents data_tensor['edited_ss_latent'] = ori_latents args = self.latents_trainer.get_inference_cond(**data_tensor) args['cond'] = torch.clone(args['edited_cond_img']) time_c = time.time() # inference latents slat_no_norm = self.latents_sampler.sample( self.latents_trainer.models['denoiser'], noise=noise, **args, steps=self.total_steps, cfg_strength=0, rescale_t=3.0, start_step=0, end_step=self.total_steps, verbose=False, ).samples time_d = time.time() if print_time: print(f'------Latents inference time: {time_d - time_c} seconds') # print(torch.mean(slat_no_norm.feats, dim=0), torch.std(slat_no_norm.feats, dim=0)) std = torch.tensor(self.ori_pipeline.slat_normalization['std'])[None].to(slat_no_norm.device) mean = torch.tensor(self.ori_pipeline.slat_normalization['mean'])[None].to(slat_no_norm.device) slat = slat_no_norm * std + mean if slat_save_path: os.makedirs(os.path.dirname(slat_save_path), exist_ok=True) np.savez_compressed( slat_save_path, feats=slat.feats.detach().cpu().numpy(), coords=slat.coords.detach().cpu().numpy(), feats_no_norm=slat_no_norm.feats.detach().cpu().numpy(), ) if not (output_video or output_mesh or render_gs_views): return output_format = ['gaussian'] if output_mesh: output_format.append('mesh') decoded_slat = self.ori_pipeline.decode_slat(slat, output_format) gaussian = decoded_slat['gaussian'] if output_video: # render the edited voxel and gaussian voxel_vis = render_image_list(voxel, resolution=vis_resolution) gaussian_vis = render_utils.render_video(gaussian[0], resolution=vis_resolution, bg_color=self.gs_bg_color, verbose=False)['color'] # render the ori voxel and gaussian ori_voxel = self.latents_trainer.dataset.decode_latent(ori_ss_latent) > 0 ori_voxel_vis = render_image_list(ori_voxel, resolution=vis_resolution) ori_gaussian = self.ori_pipeline.decode_slat(ori_latents * std + mean, ['gaussian'])['gaussian'] ori_gaussian_vis = render_utils.render_video(ori_gaussian[0], resolution=vis_resolution, bg_color=self.gs_bg_color, verbose=False)['color'] processed_ori_img = np.array(processed_ori_img.resize((vis_resolution, vis_resolution), Image.Resampling.LANCZOS)) processed_img = np.array(processed_img.resize((vis_resolution, vis_resolution), Image.Resampling.LANCZOS)) vis_images = [np.concatenate([processed_ori_img, ori_voxel_vis[i], ori_gaussian_vis[i], processed_img, voxel_vis[i], gaussian_vis[i]], axis=1) for i in range(len(voxel_vis))] imageio.mimsave(video_save_path, vis_images, fps=30) if render_gs_views: gs_dir = os.path.join(os.path.dirname(mesh_save_path), 'render_gs') _render_gaussian_eval_views( gaussian[0], gs_dir, total_render_view_num=gs_render_total_view_num, eval_view_num=gs_render_eval_view_num, resolution=gs_render_resolution, bg_color=self.gs_bg_color, verbose=False, ) if output_mesh: time_e = time.time() # export ori glb glb = postprocessing_utils.to_glb( gaussian[0], decoded_slat['mesh'][0], # Optional parameters simplify=0.95, # Ratio of triangles to remove in the simplification process texture_size=1024, # Size of the texture used for the GLB verbose=False ) glb.export(mesh_save_path) time_f = time.time() if print_time: print(f'------Mesh export time: {time_f - time_e} seconds') if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--ss_latents_config', type=str, default='ss_flow_img_dit_L_16l8_fp16.json') parser.add_argument('--ss_latents_load_id', type=str, default='img_to_voxel', help='Name under work_dirs/Editing_Training when using manual layout') parser.add_argument('--ss_latents_load_dir', type=str, default='/mnt/zsn/ckpts/3DEditFormer/img_to_voxel', help='Checkpoint root (contains ckpts/) for img_to_voxel') parser.add_argument('--ss_latents_load_ckpt', type=int, default=40000) parser.add_argument('--latents_config', type=str, default='slat_flow_img_dit_L_64l8p2_fp16.json') parser.add_argument('--latents_load_id', type=str, default='voxel_to_texture', help='Name under work_dirs/Editing_Training when using manual layout') parser.add_argument('--latents_load_dir', type=str, default='/mnt/zsn/ckpts/3DEditFormer/voxel_to_texture', help='Checkpoint root (contains ckpts/) for voxel_to_texture') parser.add_argument('--latents_load_ckpt', type=int, default=40000) parser.add_argument('--load_ema_model_for_inference', action='store_true', help='Load ema model for inference') parser.add_argument('--trellis_pipeline_path', type=str, default='/mnt/zsn/ckpts/TRELLIS-image-large', help='Local dir or repo id for TrellisImageTo3DPipeline (decode / SLAT)') parser.add_argument('--save_name', type=str, default='3DEditFormer') parser.add_argument('--blender_path', type=str, default='/opt/blender-4.2.19-linux-x64/blender') parser.add_argument('--dataset_root_dir', type=str, default='/mnt/zsn/data/3DEditVerse') parser.add_argument('--data_info_json_path', type=str, default='dataset_info.json') parser.add_argument('--select_json_path', type=str, default='test_data_info.json') parser.add_argument('--flux_edit_root_path', type=str, default='flux_edit') parser.add_argument('--alpaca_root_path', type=str, default='alpaca') parser.add_argument('--mixamo_test_animation', type=list, default=['Michelle', 'Castle Guard 01']) parser.add_argument('--mixamo_root_path', type=str, default='mixamo') parser.add_argument('--output_video', action='store_true', help='Output video') parser.add_argument('--output_mesh', action='store_true', help='Output mesh') parser.add_argument( '--save_slat', action='store_true', help='Save predicted SLAT to predict_slat.npz (feats, coords, feats_no_norm)', ) parser.add_argument( '--render_gs_views', action='store_true', help='Rasterize decoded Gaussians under render_gs/ using the same cameras as mesh Blender render', ) parser.add_argument( '--gs_render_resolution', type=int, default=512, help='Resolution for GS eval views (Blender mesh render uses 512 by default)', ) parser.add_argument('--print_time', action='store_true', help='Print time') parser.add_argument('--total_render_view_num', type=int, default=300) parser.add_argument('--eval_view_num', type=int, default=10) parser.add_argument('--engine', type=str, default='CYCLES', choices=['CYCLES', 'BLENDER_EEVEE_NEXT']) parser.add_argument('--BLENDER_EEVEE_NEXT_on_0_device', action='store_true', help='BLENDER_EEVEE_NEXT on 0 device') parser.add_argument('--cuda_idx', type=int, nargs='*', default=[0]) parser.add_argument('--world_size', type=int, default=1, help='Total number of GPUs') parser.add_argument('--rank', type=int, default=0, help='Current GPU rank (0 to world_size-1)') parser.add_argument('--start_idx', type=int, default=None, help='Start index') parser.add_argument('--end_idx', type=int, default=None, help='End index') parser.add_argument('--render_timeout_seconds', type=int, default=100, help='Render timeout in seconds') parser.add_argument('--debug', action='store_true', help='Debug') parser.add_argument('--empty_structure_fallback', type=str, default='error', choices=['error', 'original'], help='How to handle an empty sparse structure prediction') parser.add_argument( '--skip_blender_render', action='store_true', help='Skip Blender multi-view rendering queue and exit after mesh inference (renders can be filled later)', ) parser.add_argument( '--fixed_edited_ss_predictions_jsonl', type=str, default='', help=( 'If set, JSONL with sample_id and pred_ss_path (e.g. stage1 eval output); ' 'when sample_id matches eval key, skip SS sampling and load that npz mean.' ), ) args = parser.parse_args() args.data_info_json_path = os.path.join(args.dataset_root_dir, args.data_info_json_path) args.select_json_path = os.path.join(args.dataset_root_dir, args.select_json_path) args.flux_edit_root_path = os.path.join(args.dataset_root_dir, args.flux_edit_root_path) args.alpaca_root_path = os.path.join(args.dataset_root_dir, args.alpaca_root_path) args.mixamo_root_path = os.path.join(args.dataset_root_dir, args.mixamo_root_path) set_seed(42) assert len(args.cuda_idx) == 1 # load testing data save_path = f'./work_dirs/eval_results/{args.save_name}' with open(args.select_json_path, 'r') as f: select_data = json.load(f) with open(args.data_info_json_path, 'r') as f: data_info = json.load(f) fixed_edited_ss_by_key = {} if getattr(args, 'fixed_edited_ss_predictions_jsonl', '') and args.fixed_edited_ss_predictions_jsonl: pred_path = args.fixed_edited_ss_predictions_jsonl with open(pred_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue row = json.loads(line) fixed_edited_ss_by_key[row['sample_id']] = row['pred_ss_path'] processed_keys = [] for key, values in select_data.items(): for value in values: processed_keys.append((key, value)) print('processed keys: ', len(processed_keys)) if args.start_idx is not None and args.end_idx is not None: start_idx = args.start_idx end_idx = args.end_idx processed_keys = processed_keys[start_idx: end_idx] else: chunk_size = math.ceil(len(processed_keys) / args.world_size) start_idx = args.rank * chunk_size end_idx = min((args.rank + 1) * chunk_size, len(processed_keys)) processed_keys = processed_keys[start_idx: end_idx] print(f'World size: {args.world_size}, Rank {args.rank}, tasks index: {start_idx} - {end_idx}') if args.debug: processed_keys = processed_keys[:6] trellis_edititng_model = TrellisEdititngModel( ss_latents_config_path=f'./configs/editing/{args.ss_latents_config}', ss_latents_load_dir=args.ss_latents_load_dir, ss_latents_load_ckpt=args.ss_latents_load_ckpt, latents_config_path=f'./configs/editing/{args.latents_config}', latents_load_dir=args.latents_load_dir, latents_load_ckpt=args.latents_load_ckpt, load_ema_model_for_inference=args.load_ema_model_for_inference, trellis_pipeline_path=args.trellis_pipeline_path, ) print('init trellis edititng model done') if args.skip_blender_render: print('skip_blender_render=True: Blender render worker disabled') render_queue = queue.Queue() def render_worker(): while True: try: task = render_queue.get(timeout=1) if task is None: break mesh_save_path, render_args = task output_folder = _render( file_path=mesh_save_path, engine=render_args['engine'], cuda_idx=render_args['cuda_idx'], total_render_view_num=render_args['total_render_view_num'], eval_view_num=render_args['eval_view_num'], BLENDER_EEVEE_NEXT_on_0_device=render_args['BLENDER_EEVEE_NEXT_on_0_device'], timeout_seconds=render_args['timeout_seconds'], no_norm_scene=True, debug=render_args['debug'], blender_path=render_args['blender_path'], ) print(f'render done: {mesh_save_path}') render_queue.task_done() except queue.Empty: continue except Exception as e: print(f'render error: {e}') render_queue.task_done() render_thread = None if not args.skip_blender_render: render_thread = threading.Thread(target=render_worker, daemon=True) render_thread.start() for data in tqdm(processed_keys): dataset_type = data[0] if dataset_type == 'alpaca': key = data[1] ori_ss_latents_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]['ori_ss_latents_path']) ori_latents_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]['ori_latents_path']) ori_img_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]['ori_img_path']) edit_img_path = os.path.join(args.alpaca_root_path, data_info[dataset_type][key]['edit_img_path']) elif dataset_type == 'flux_edit': key = data[1] ori_ss_latents_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]['ori_ss_latents_path']) ori_latents_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]['ori_latents_path']) ori_img_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]['ori_img_path']) edit_img_path = os.path.join(args.flux_edit_root_path, data_info[dataset_type][key]['edit_img_path']) elif dataset_type == 'mixamo': character_name, ori_idx, edit_idx = data[1] key = f'{character_name}_{ori_idx}_{edit_idx}' ori_ss_latents_path = os.path.join(args.mixamo_root_path, data_info[dataset_type][character_name][ori_idx]['ss_latents_path']) ori_latents_path = os.path.join(args.mixamo_root_path, data_info[dataset_type][character_name][ori_idx]['latents_path']) ori_img_path = os.path.join(args.mixamo_root_path, data_info[dataset_type][character_name][ori_idx]['img_path']) edit_img_path = os.path.join(args.mixamo_root_path, data_info[dataset_type][character_name][edit_idx]['img_path']) else: raise ValueError(f'Invalid dataset type: {dataset_type}') video_save_path = os.path.join(save_path, dataset_type, key, 'video.mp4') mesh_save_path = os.path.join(save_path, dataset_type, key, 'edit.glb') slat_save_path = os.path.join(save_path, dataset_type, key, 'predict_slat.npz') os.makedirs(os.path.dirname(video_save_path), exist_ok=True) gs_render_marker = os.path.join(os.path.dirname(mesh_save_path), 'render_gs', 'transforms.json') need_run = not os.path.exists(mesh_save_path) if args.save_slat and not os.path.exists(slat_save_path): need_run = True if args.render_gs_views and not os.path.exists(gs_render_marker): need_run = True if need_run: time_a = time.time() trellis_edititng_model.editing_inference( ori_ss_latents_path=ori_ss_latents_path, ori_latents_path=ori_latents_path, edited_img_path=edit_img_path, ori_img_path=ori_img_path, output_video=args.output_video, output_mesh=args.output_mesh, video_save_path=video_save_path, mesh_save_path=mesh_save_path, slat_save_path=slat_save_path if args.save_slat else None, render_gs_views=args.render_gs_views, gs_render_total_view_num=args.total_render_view_num, gs_render_eval_view_num=args.eval_view_num, gs_render_resolution=args.gs_render_resolution, ori_latents_norm=True if dataset_type == 'mixamo' else False, print_time=args.print_time, empty_structure_fallback=args.empty_structure_fallback, edited_ss_latents_path=fixed_edited_ss_by_key.get(key), ) time_b = time.time() if args.print_time: print(f'--------- Total Editing inference time: {time_b - time_a} seconds') if (not args.skip_blender_render) and args.output_mesh and os.path.exists(mesh_save_path) and ( not os.path.exists(os.path.join(os.path.dirname(mesh_save_path), 'render', 'transforms.json')) ): render_args = { 'engine': args.engine, 'cuda_idx': args.cuda_idx[0], 'total_render_view_num': args.total_render_view_num, 'eval_view_num': args.eval_view_num, 'BLENDER_EEVEE_NEXT_on_0_device': args.BLENDER_EEVEE_NEXT_on_0_device, 'timeout_seconds': args.render_timeout_seconds, 'debug': args.debug, 'blender_path': args.blender_path, } render_queue.put((mesh_save_path, render_args)) print(f'render task submitted: {mesh_save_path}') if args.skip_blender_render: print('skip_blender_render=True: skipping Blender queue drain') else: print('waiting for all render tasks to finish...') render_queue.join() render_queue.put(None) render_thread.join() print('all render tasks done')