Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,046 Bytes
3075458 |
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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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()
|