cf3b919b65
* feat(sim): Add auto scale in convex decomposition. * feat(texture): Optimize back-projected texture quality. * feat(texture): Add `texture-cli`.
124 lines
4.0 KiB
Python
124 lines
4.0 KiB
Python
import os
|
|
import shutil
|
|
from dataclasses import dataclass
|
|
|
|
import tyro
|
|
from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
|
|
from embodied_gen.data.differentiable_render import entrypoint as drender_api
|
|
from embodied_gen.data.utils import as_list
|
|
from embodied_gen.models.delight_model import DelightingModel
|
|
from embodied_gen.models.sr_model import ImageRealESRGAN
|
|
from embodied_gen.scripts.render_mv import (
|
|
build_texture_gen_pipe,
|
|
)
|
|
from embodied_gen.scripts.render_mv import infer_pipe as render_mv_api
|
|
from embodied_gen.utils.log import logger
|
|
|
|
|
|
@dataclass
|
|
class TextureGenConfig:
|
|
mesh_path: str | list[str]
|
|
prompt: str | list[str]
|
|
output_root: str
|
|
controlnet_cond_scale: float = 0.7
|
|
guidance_scale: float = 9
|
|
strength: float = 0.9
|
|
num_inference_steps: int = 40
|
|
delight: bool = True
|
|
seed: int = 0
|
|
base_ckpt_dir: str = "./weights"
|
|
texture_size: int = 2048
|
|
ip_adapt_scale: float = 0.0
|
|
ip_img_path: str | list[str] | None = None
|
|
|
|
|
|
def entrypoint() -> None:
|
|
cfg = tyro.cli(TextureGenConfig)
|
|
cfg.mesh_path = as_list(cfg.mesh_path)
|
|
cfg.prompt = as_list(cfg.prompt)
|
|
cfg.ip_img_path = as_list(cfg.ip_img_path)
|
|
assert len(cfg.mesh_path) == len(cfg.prompt)
|
|
|
|
# Pre-load models.
|
|
if cfg.ip_adapt_scale > 0:
|
|
PIPELINE = build_texture_gen_pipe(
|
|
base_ckpt_dir="./weights",
|
|
ip_adapt_scale=cfg.ip_adapt_scale,
|
|
device="cuda",
|
|
)
|
|
else:
|
|
PIPELINE = build_texture_gen_pipe(
|
|
base_ckpt_dir="./weights",
|
|
ip_adapt_scale=0,
|
|
device="cuda",
|
|
)
|
|
DELIGHT = None
|
|
if cfg.delight:
|
|
DELIGHT = DelightingModel()
|
|
IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
|
|
|
|
for idx in range(len(cfg.mesh_path)):
|
|
mesh_path = cfg.mesh_path[idx]
|
|
prompt = cfg.prompt[idx]
|
|
uuid = os.path.splitext(os.path.basename(mesh_path))[0]
|
|
output_root = os.path.join(cfg.output_root, uuid)
|
|
drender_api(
|
|
mesh_path=mesh_path,
|
|
output_root=f"{output_root}/condition",
|
|
uuid=uuid,
|
|
)
|
|
render_mv_api(
|
|
index_file=f"{output_root}/condition/index.json",
|
|
controlnet_cond_scale=cfg.controlnet_cond_scale,
|
|
guidance_scale=cfg.guidance_scale,
|
|
strength=cfg.strength,
|
|
num_inference_steps=cfg.num_inference_steps,
|
|
ip_adapt_scale=cfg.ip_adapt_scale,
|
|
ip_img_path=(
|
|
None if cfg.ip_img_path is None else cfg.ip_img_path[idx]
|
|
),
|
|
prompt=prompt,
|
|
save_dir=f"{output_root}/multi_view",
|
|
sub_idxs=[[0, 1, 2], [3, 4, 5]],
|
|
pipeline=PIPELINE,
|
|
seed=cfg.seed,
|
|
)
|
|
textured_mesh = backproject_api(
|
|
delight_model=DELIGHT,
|
|
imagesr_model=IMAGESR_MODEL,
|
|
mesh_path=mesh_path,
|
|
color_path=f"{output_root}/multi_view/color_sample0.png",
|
|
output_path=f"{output_root}/texture_mesh/{uuid}.obj",
|
|
save_glb_path=f"{output_root}/texture_mesh/{uuid}.glb",
|
|
skip_fix_mesh=True,
|
|
delight=cfg.delight,
|
|
no_save_delight_img=True,
|
|
texture_wh=[cfg.texture_size, cfg.texture_size],
|
|
)
|
|
drender_api(
|
|
mesh_path=f"{output_root}/texture_mesh/{uuid}.obj",
|
|
output_root=f"{output_root}/texture_mesh",
|
|
uuid=uuid,
|
|
num_images=90,
|
|
elevation=[20],
|
|
with_mtl=True,
|
|
gen_color_mp4=True,
|
|
pbr_light_factor=1.2,
|
|
)
|
|
|
|
# Re-organize folders
|
|
shutil.rmtree(f"{output_root}/condition")
|
|
shutil.copy(
|
|
f"{output_root}/texture_mesh/{uuid}/color.mp4",
|
|
f"{output_root}/color.mp4",
|
|
)
|
|
shutil.rmtree(f"{output_root}/texture_mesh/{uuid}")
|
|
|
|
logger.info(
|
|
f"Successfully generate textured mesh in {output_root}/texture_mesh"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
entrypoint()
|