Spaces:
Running
on
Zero
Running
on
Zero
| # Open Source Model Licensed under the Apache License Version 2.0 | |
| # and Other Licenses of the Third-Party Components therein: | |
| # The below Model in this distribution may have been modified by THL A29 Limited | |
| # ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited. | |
| # Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved. | |
| # The below software and/or models in this distribution may have been | |
| # modified by THL A29 Limited ("Tencent Modifications"). | |
| # All Tencent Modifications are Copyright (C) THL A29 Limited. | |
| # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT | |
| # except for the third-party components listed below. | |
| # Hunyuan 3D does not impose any additional limitations beyond what is outlined | |
| # in the repsective licenses of these third-party components. | |
| # Users must comply with all terms and conditions of original licenses of these third-party | |
| # components and must ensure that the usage of the third party components adheres to | |
| # all relevant laws and regulations. | |
| # For avoidance of doubts, Hunyuan 3D means the large language models and | |
| # their software and algorithms, including trained model weights, parameters (including | |
| # optimizer states), machine-learning model code, inference-enabling code, training-enabling code, | |
| # fine-tuning enabling code and other elements of the foregoing made publicly available | |
| # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT. | |
| import os, sys | |
| sys.path.insert(0, f"{os.path.dirname(os.path.dirname(os.path.abspath(__file__)))}") | |
| import time | |
| import torch | |
| import random | |
| import numpy as np | |
| from PIL import Image | |
| from einops import rearrange | |
| from PIL import Image, ImageSequence | |
| from infer.utils import seed_everything, timing_decorator, auto_amp_inference | |
| from infer.utils import get_parameter_number, set_parameter_grad_false, str_to_bool | |
| from svrm.predictor import MV23DPredictor | |
| class Views2Mesh(): | |
| def __init__(self, mv23d_cfg_path, mv23d_ckt_path, | |
| device="cuda:0", use_lite=False, save_memory=False): | |
| ''' | |
| mv23d_cfg_path: config yaml file | |
| mv23d_ckt_path: path to ckpt | |
| use_lite: lite version | |
| save_memory: cpu auto | |
| ''' | |
| self.mv23d_predictor = MV23DPredictor(mv23d_ckt_path, mv23d_cfg_path, device=device) | |
| self.mv23d_predictor.model.eval() | |
| self.order = [0, 1, 2, 3, 4, 5] if use_lite else [0, 2, 4, 5, 3, 1] | |
| self.device = device | |
| self.save_memory = save_memory | |
| set_parameter_grad_false(self.mv23d_predictor.model) | |
| print('view2mesh model', get_parameter_number(self.mv23d_predictor.model)) | |
| def __call__(self, *args, **kwargs): | |
| if self.save_memory: | |
| self.mv23d_predictor.model = self.mv23d_predictor.model.to(self.device) | |
| torch.cuda.empty_cache() | |
| res = self.call(*args, **kwargs) | |
| self.mv23d_predictor.model = self.mv23d_predictor.model.to("cpu") | |
| else: | |
| res = self.call(*args, **kwargs) | |
| torch.cuda.empty_cache() | |
| return res | |
| def call( | |
| self, | |
| views_pil=None, | |
| cond_pil=None, | |
| gif_pil=None, | |
| seed=0, | |
| target_face_count = 10000, | |
| do_texture_mapping = True, | |
| save_folder='./outputs/test' | |
| ): | |
| ''' | |
| can set views_pil, cond_pil simutaously or set gif_pil only | |
| seed: int | |
| target_face_count: int | |
| save_folder: path to save mesh files | |
| ''' | |
| save_dir = save_folder | |
| os.makedirs(save_dir, exist_ok=True) | |
| if views_pil is not None and cond_pil is not None: | |
| show_image = rearrange(np.asarray(views_pil, dtype=np.uint8), | |
| '(n h) (m w) c -> (n m) h w c', n=3, m=2) | |
| views = [Image.fromarray(show_image[idx]) for idx in self.order] | |
| image_list = [cond_pil]+ views | |
| image_list = [img.convert('RGB') for img in image_list] | |
| elif gif_pil is not None: | |
| image_list = [img.convert('RGB') for img in ImageSequence.Iterator(gif_pil)] | |
| image_input = image_list[0] | |
| image_list = image_list[1:] + image_list[:1] | |
| seed_everything(seed) | |
| self.mv23d_predictor.predict( | |
| image_list, | |
| save_dir = save_dir, | |
| image_input = image_input, | |
| target_face_count = target_face_count, | |
| do_texture_mapping = do_texture_mapping | |
| ) | |
| torch.cuda.empty_cache() | |
| return save_dir | |
| if __name__ == "__main__": | |
| import argparse | |
| def get_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--views_path", type=str, required=True) | |
| parser.add_argument("--cond_path", type=str, required=True) | |
| parser.add_argument("--save_folder", default="./outputs/test/", type=str) | |
| parser.add_argument("--mv23d_cfg_path", default="./svrm/configs/svrm.yaml", type=str) | |
| parser.add_argument("--mv23d_ckt_path", default="weights/svrm/svrm.safetensors", type=str) | |
| parser.add_argument("--max_faces_num", default=90000, type=int, | |
| help="max num of face, suggest 90000 for effect, 10000 for speed") | |
| parser.add_argument("--device", default="cuda:0", type=str) | |
| parser.add_argument("--use_lite", default='false', type=str) | |
| parser.add_argument("--do_texture_mapping", default='false', type=str) | |
| return parser.parse_args() | |
| args = get_args() | |
| args.use_lite = str_to_bool(args.use_lite) | |
| args.do_texture_mapping = str_to_bool(args.do_texture_mapping) | |
| views = Image.open(args.views_path) | |
| cond = Image.open(args.cond_path) | |
| views_to_mesh_model = Views2Mesh( | |
| args.mv23d_cfg_path, | |
| args.mv23d_ckt_path, | |
| device = args.device, | |
| use_lite = args.use_lite | |
| ) | |
| views_to_mesh_model( | |
| views, cond, 0, | |
| target_face_count = args.max_faces_num, | |
| save_folder = args.save_folder, | |
| do_texture_mapping = args.do_texture_mapping | |
| ) | |