diff --git a/Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_texture.py b/Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_texture.py new file mode 100644 index 0000000..a177eb9 --- /dev/null +++ b/Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_texture.py @@ -0,0 +1,452 @@ +import os +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from spandrel import ModelLoader + +from MV_Adapter.mvadapter.utils import image_to_tensor, make_image_grid, tensor_to_image +from MV_Adapter.mvadapter.utils.mesh_utils import ( + Camera, + CameraProjection, + NVDiffRastContextWrapper, + SmartPainter, + TexturedMesh, + get_camera, + get_orthogonal_camera, + load_mesh, + render, + replace_mesh_texture_and_save, +) +from MV_Adapter.mvadapter.utils.mesh_utils.mesh_process import process_raw + +# Import dilate function to remove texture seams +def torch_image_dilate(img): + """ + Remove thin seams on generated texture + img (torch): [H, W, C] + """ + import cv2 + img = np.asarray(img.cpu().numpy(), dtype=np.float32) + img = img * 255 + img = img.clip(0, 255) + mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True) + mask = (mask <= 3.0).astype(np.float32) + kernel = np.ones((3, 3), 'uint8') + dilate_img = cv2.dilate(img, kernel, iterations=1) + img = img * (1 - mask) + dilate_img * mask + img = (img.clip(0, 255) / 255).astype(np.float32) + return torch.from_numpy(img) + + +def clear(): + torch.cuda.empty_cache() + + +@contextmanager +def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor): + texture_ = mesh.texture + mesh.texture = texture + try: + yield + finally: + mesh.texture = texture_ + + +@dataclass +class ModProcessConfig: + view_upscale: bool = False + view_upscale_factor: int = 2 + inpaint_mode: str = "uv" # in ["none", "uv", "view"] + view_inpaint_max_view_score_thresh: float = 0.02 + view_inpaint_min_rounds: int = 4 + view_inpaint_max_rounds: int = 8 + view_inpaint_uv_padding_end: bool = True + + +@dataclass +class TexturePipelineOutput: + shaded_model_save_path: Optional[str] = None + pbr_model_save_path: Optional[str] = None + + +class TexturePipeline: + def __init__(self, upscaler_ckpt_path: str, inpaint_ckpt_path: str, device: str): + self.device = device + self.ctx = NVDiffRastContextWrapper(device=self.device) + self.cam_proj = CameraProjection( + pb_backend="torch-cuda", bg_remover=None, device=self.device + ) + if upscaler_ckpt_path is not None: + self.upscaler = ModelLoader().load_from_file(upscaler_ckpt_path) + self.upscaler.to(self.device).eval().half() + if inpaint_ckpt_path is not None: + self.inpainter = ModelLoader().load_from_file(inpaint_ckpt_path) + self.inpainter.to(self.device).eval() + + self.smart_painter = SmartPainter(self.device) + + def load_packed_images(self, packed_image_path: Optional[str]) -> List[Image.Image]: + if packed_image_path is None: + return None + packed_image = Image.open(packed_image_path) + images = np.array_split(np.array(packed_image), 6, axis=1) + images = [Image.fromarray(im) for im in images] + return images + + def maybe_upscale_image( + self, + tensor: Optional[torch.FloatTensor], + upscale: bool, + upscale_factor: int, + batched: bool = False, + ) -> Optional[torch.FloatTensor]: + if upscale: + with torch.no_grad(): + tensor = tensor.permute(0, 3, 1, 2) + if batched: + tensor = self.upscaler(tensor.half()).float() + else: + tensor = torch.concat( + [ + self.upscaler(im.unsqueeze(0).half()).float() + for im in tensor + ], + dim=0, + ) + tensor = tensor.clamp(0, 1).permute(0, 2, 3, 1) + clear() + return tensor + + def view_inpaint( + self, + mod_name: str, + mesh: TexturedMesh, + uv_proj: torch.FloatTensor, + uv_valid_mask: torch.BoolTensor, + config: ModProcessConfig, + debug_dir: Optional[str] = None, + ) -> Tuple[torch.FloatTensor, torch.BoolTensor]: + def inpaint_func( + image: torch.FloatTensor, mask: torch.FloatTensor + ) -> torch.FloatTensor: + with torch.no_grad(): + return self.inpainter(image.to(torch.float32), mask.to(torch.float32)) + + return self.smart_painter( + mod_name, + mesh, + inpaint_func, + uv_proj, + ~uv_valid_mask, + max_view_score_thresh=config.view_inpaint_max_view_score_thresh, + min_rounds=config.view_inpaint_min_rounds, + max_rounds=config.view_inpaint_max_rounds, + uv_padding_end=config.view_inpaint_uv_padding_end, + debug_dir=debug_dir, + debug_visualize_details=False, + ) + + def __call__( + self, + mesh_path: str, + save_dir: str, + save_name: str = "default", + # mesh load settings + move_to_center: bool = False, + front_x: bool = True, + # uv unwarp + uv_unwarp: bool = False, + preprocess_mesh: bool = False, + # projection + uv_size: int = 4096, + # modes + rgb_path: Optional[str] = None, + rgb_process_config: ModProcessConfig = ModProcessConfig(), + base_color_path: Optional[str] = None, + base_color_process_config: ModProcessConfig = ModProcessConfig(), + orm_path: Optional[str] = None, + orm_process_config: ModProcessConfig = ModProcessConfig(), + normal_path: Optional[str] = None, + normal_strength: float = 1.0, + normal_process_config: ModProcessConfig = ModProcessConfig(), + # inpaint + uv_inpaint_use_network: bool = False, + view_inpaint_include_occlusion_boundary: bool = False, + poisson_reprojection: bool = False, + # camera + camera_projection_type: str = "ORTHO", + camera_elevation_deg: List[float] = [0, 0, 0, 0, 89.99, -89.99], + camera_azimuth_deg: List[float] = [0, 90, 180, 270, 180, 180], + camera_distance: float = 1.0, + camera_ortho_scale: float = 1.1, + camera_fov_deg: float = 40, + # debug + debug_mode: bool = False, + # seam removal + apply_dilate: bool = True, + ): + clear() + + if debug_mode: + debug_dir = os.path.join(save_dir, "debug") + os.makedirs(debug_dir, exist_ok=True) + + if uv_unwarp: + file_suffix = os.path.splitext(mesh_path)[-1] + mesh_path_new = mesh_path.replace(file_suffix, f"_unwarp{file_suffix}") + process_raw(mesh_path, mesh_path_new, preprocess=preprocess_mesh) + mesh_path = mesh_path_new + + mesh: TexturedMesh = load_mesh( + mesh_path, + rescale=True, + move_to_center=move_to_center, + front_x_to_y=front_x, + default_uv_size=uv_size, + device=self.device, + ) + + # projection + if camera_projection_type == "PERSP": + raise NotImplementedError + elif camera_projection_type == "ORTHO": + cameras = get_orthogonal_camera( + elevation_deg=camera_elevation_deg, + distance=[camera_distance] * 6, + left=-camera_ortho_scale / 2, + right=camera_ortho_scale / 2, + bottom=-camera_ortho_scale / 2, + top=camera_ortho_scale / 2, + azimuth_deg=[x - 90 for x in camera_azimuth_deg], # -y as front + device=self.device, + ) + + mod_kwargs = { + "rgb": (rgb_path, rgb_process_config), + "base_color": (base_color_path, base_color_process_config), + "orm": (orm_path, orm_process_config), + "normal": (normal_path, normal_process_config), + } + mod_uv_image, mod_uv_tensor = {}, {} + for mod_name, (mod_path, mod_process_config) in mod_kwargs.items(): + if mod_path is None: + mod_uv_image[mod_name] = None + mod_uv_tensor[mod_name] = None + continue + mod_images = self.load_packed_images(mod_path) + mod_tensor = image_to_tensor(mod_images, device=self.device) + mod_tensor = self.maybe_upscale_image( + mod_tensor, + mod_process_config.view_upscale, + mod_process_config.view_upscale_factor, + ) + if mod_process_config.view_upscale and debug_mode: + make_image_grid(tensor_to_image(mod_tensor, batched=True), rows=1).save( + os.path.join(debug_dir, f"{mod_name}_upscaled.jpg") + ) + + if mod_name == "normal": + _, height, width, _ = mod_tensor.shape + render_out = render( + self.ctx, + mesh, + cameras, + height=height, + width=width, + render_attr=False, + render_depth=False, + render_normal=True, + render_tangent=True, + ) + + # compute UV tangent space + vN = render_out.normal + vT = render_out.tangent + vB = torch.cross(vN, vT, dim=-1) + tang_space = F.normalize(torch.stack([vT, vB, vN], dim=-2), dim=-1) + + # compute geometry tangent space + vGN = vN + vGT = torch.as_tensor( + [ + [1, 0, 0], + [0, 1, 0], + [-1, 0, 0], + [0, -1, 0], + [-1, 0, 0], + [-1, 0, 0], + ], + dtype=torch.float32, + device=self.device, + )[:, None, None, :] + vGB = torch.cross(vGN, vGT, dim=-1) + vGT = torch.cross(vGB, vGN, dim=-1) + geo_tang_space = F.normalize( + torch.stack([vGT, vGB, vGN], dim=-2), dim=-1 + ) + + # restore world-space normal from geometry tangent space + mod_tensor = mod_tensor * 2 - 1 + mod_tensor = F.normalize( + ( + mod_tensor[:, :, :, None, :] + * geo_tang_space.permute(0, 1, 2, 4, 3) + ).sum(-1), + dim=-1, + ) + + # bake world-space normal to UV tangent space + mod_tensor = F.normalize( + (mod_tensor[:, :, :, None, :] * tang_space).sum(-1), dim=-1 + ) + mod_tensor = (mod_tensor * 0.5 + 0.5).clamp(0, 1) + + uv_proj, uv_valid_mask = self.cam_proj( + mod_tensor, + mesh, + cameras, + from_scratch=mod_process_config.inpaint_mode != "none", + poisson_blending=False, + depth_grad_dilation=5, + uv_exp_blend_alpha=3, + uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]), + aoi_cos_valid_threshold=0.2, + uv_size=uv_size, + return_uv_projection_mask=True, + ) + uv_proj[~uv_valid_mask] = torch.as_tensor([0.5, 0.5, 1]).to(uv_proj) + else: + # TODO: tweak depth_grad_dilation + cam_proj_out = self.cam_proj( + mod_tensor, + mesh, + cameras, + from_scratch=mod_process_config.inpaint_mode != "none", + poisson_blending=False, + depth_grad_dilation=5, + depth_grad_threshold=0.1, + uv_exp_blend_alpha=3, + uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]), + aoi_cos_valid_threshold=0.2, + uv_size=uv_size, + uv_padding=not uv_inpaint_use_network, + return_dict=True, + ) + uv_proj, uv_valid_mask, uv_depth_grad = ( + cam_proj_out.uv_proj, + cam_proj_out.uv_proj_mask, + cam_proj_out.uv_depth_grad, + ) + if uv_inpaint_use_network: + uv_inpaint_mask_input = 1 - uv_valid_mask[None, None].float() + uv_inpaint_image_input = uv_proj[None].permute(0, 3, 1, 2) + with torch.no_grad(): + uv_inpaint_result = self.inpainter( + uv_inpaint_image_input, uv_inpaint_mask_input + )[0].permute(1, 2, 0) + clear() + if debug_mode: + make_image_grid( + [ + tensor_to_image(uv_proj), + tensor_to_image(uv_valid_mask), + tensor_to_image(uv_inpaint_result), + ] + ).save(os.path.join(debug_dir, f"{mod_name}_uv_inpaint.jpg")) + uv_proj = uv_inpaint_result.contiguous() + + if mod_process_config.inpaint_mode == "view": + if view_inpaint_include_occlusion_boundary: + uv_max_depth_grad = uv_depth_grad.max(dim=0)[0] + uv_valid_mask = uv_valid_mask & (uv_max_depth_grad < 0.1) + uv_proj, uv_valid_mask = self.view_inpaint( + mod_name, + mesh, + uv_proj, + uv_valid_mask, + mod_process_config, + debug_dir=debug_dir if debug_mode else None, + ) + + if poisson_reprojection: + # up and down + mesh.texture = uv_proj + uv_proj = self.cam_proj( + mod_tensor[4:5], + mesh, + cameras[4:5], + from_scratch=False, + poisson_blending=True, + pb_keep_original_border=True, + depth_grad_dilation=5, + uv_exp_blend_alpha=3, + uv_exp_blend_view_weight=torch.as_tensor([1, 1]), + aoi_cos_valid_threshold=0.2, + uv_size=uv_size, + uv_padding=True, + return_dict=False, + ) + # front, sides and back + mesh.texture = uv_proj + uv_proj = self.cam_proj( + mod_tensor[0:4], + mesh, + cameras[0:4], + from_scratch=False, + poisson_blending=True, + pb_keep_original_border=True, + depth_grad_dilation=5, + uv_exp_blend_alpha=3, + uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1]), + aoi_cos_valid_threshold=0.2, + uv_size=uv_size, + uv_padding=True, + return_dict=False, + ) + + if mod_name == "orm": + uv_proj[:, :, 0] = 1.0 + + # Apply dilate to remove texture seams before final conversion + if mod_name in ["rgb", "base_color"] and apply_dilate: # Apply to color textures only + uv_proj = torch_image_dilate(uv_proj) + + mod_uv_image[mod_name] = tensor_to_image(uv_proj) + mod_uv_tensor[mod_name] = uv_proj + clear() + + shaded_model_save_path = None + if mod_uv_image["rgb"] is not None: + shaded_model_save_path = os.path.join(save_dir, f"{save_name}_shaded.glb") + replace_mesh_texture_and_save( + mesh_path, + shaded_model_save_path, + texture=mod_uv_image["rgb"], + backend="trimesh", + task_id=save_name, + ) + pbr_model_save_path = None + if mod_uv_image["base_color"] is not None: + pbr_model_save_path = os.path.join(save_dir, f"{save_name}_pbr.glb") + replace_mesh_texture_and_save( + mesh_path, + pbr_model_save_path, + texture=mod_uv_image["base_color"], + metallic_roughness_texture=mod_uv_image["orm"], + normal_texture=mod_uv_image["normal"], + normal_strength=normal_strength, + backend="trimesh", + task_id=save_name, + ) + + clear() + + return TexturePipelineOutput( + shaded_model_save_path=shaded_model_save_path, + pbr_model_save_path=pbr_model_save_path, + ) \ No newline at end of file diff --git a/Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py b/Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py index 4aef3ec..8b4efde 100644 --- a/Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py +++ b/Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py @@ -409,80 +409,80 @@ def replace_mesh_texture_and_save_trimesh( tmesh.export(output_path, tree_postprocessor=p_func) -def replace_mesh_texture_and_save_gltflib( - input_path: str, - output_path: str, - texture: Union[Image.Image, np.ndarray, torch.Tensor], - metallic_roughness_texture: Optional[ - Union[Image.Image, np.ndarray, torch.Tensor] - ] = None, - normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None, - texture_format: str = "JPEG", - normal_strength: float = 1.0, - task_id: str = "", - **kwargs, -) -> None: - from gltflib import GLTF - from gltflib import Image as gltfImage - from gltflib import NormalTextureInfo, Texture, TextureInfo +# def replace_mesh_texture_and_save_gltflib( +# input_path: str, +# output_path: str, +# texture: Union[Image.Image, np.ndarray, torch.Tensor], +# metallic_roughness_texture: Optional[ +# Union[Image.Image, np.ndarray, torch.Tensor] +# ] = None, +# normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None, +# texture_format: str = "JPEG", +# normal_strength: float = 1.0, +# task_id: str = "", +# **kwargs, +# ) -> None: +# from gltflib import GLTF +# from gltflib import Image as gltfImage +# from gltflib import NormalTextureInfo, Texture, TextureInfo - def convert_image(tensor): - texture = tensor_to_image(tensor) - texture.format = texture_format - return texture +# def convert_image(tensor): +# texture = tensor_to_image(tensor) +# texture.format = texture_format +# return texture - def add_texture(gltf: GLTF, image: Image.Image) -> int: - jpeg_buffer = io.BytesIO() - image.save(jpeg_buffer, format=texture_format) - new_texture_data = bytearray(jpeg_buffer.getvalue()) - _, offset, bytelen = gltf._create_or_extend_glb_resource(new_texture_data) - bufferView = gltf._create_embedded_image_buffer_view(offset, bytelen) - image = gltfImage(mimeType="image/jpeg", bufferView=bufferView) - if gltf.model.images: - gltf.model.images.append(image) - else: - gltf.model.images = [image] - texture = Texture(source=len(gltf.model.images) - 1) - if gltf.model.textures: - gltf.model.textures.append(texture) - else: - gltf.model.textures = [texture] - return len(gltf.model.textures) - 1 +# def add_texture(gltf: GLTF, image: Image.Image) -> int: +# jpeg_buffer = io.BytesIO() +# image.save(jpeg_buffer, format=texture_format) +# new_texture_data = bytearray(jpeg_buffer.getvalue()) +# _, offset, bytelen = gltf._create_or_extend_glb_resource(new_texture_data) +# bufferView = gltf._create_embedded_image_buffer_view(offset, bytelen) +# image = gltfImage(mimeType="image/jpeg", bufferView=bufferView) +# if gltf.model.images: +# gltf.model.images.append(image) +# else: +# gltf.model.images = [image] +# texture = Texture(source=len(gltf.model.images) - 1) +# if gltf.model.textures: +# gltf.model.textures.append(texture) +# else: +# gltf.model.textures = [texture] +# return len(gltf.model.textures) - 1 - gltf = GLTF.load(input_path) +# gltf = GLTF.load(input_path) - texture = convert_image(texture) - if metallic_roughness_texture is not None: - metallic_roughness_texture = convert_image(metallic_roughness_texture) +# texture = convert_image(texture) +# if metallic_roughness_texture is not None: +# metallic_roughness_texture = convert_image(metallic_roughness_texture) - if normal_texture is not None: - normal_texture = convert_image(normal_texture) +# if normal_texture is not None: +# normal_texture = convert_image(normal_texture) - pbr_material = gltf.model.materials[0].pbrMetallicRoughness +# pbr_material = gltf.model.materials[0].pbrMetallicRoughness - pbr_material.baseColorTexture = TextureInfo(index=add_texture(gltf, texture)) - pbr_material.baseColorFactor = [1, 1, 1, 1] +# pbr_material.baseColorTexture = TextureInfo(index=add_texture(gltf, texture)) +# pbr_material.baseColorFactor = [1, 1, 1, 1] - if metallic_roughness_texture is not None: - pbr_material.metallicRoughnessTexture = TextureInfo( - index=add_texture(gltf, metallic_roughness_texture) - ) - pbr_material.metallicFactor = 1.0 - pbr_material.roughnessFactor = 1.0 - else: - pbr_material.metallicFactor = 0.0 - pbr_material.roughnessFactor = 0.9 +# if metallic_roughness_texture is not None: +# pbr_material.metallicRoughnessTexture = TextureInfo( +# index=add_texture(gltf, metallic_roughness_texture) +# ) +# pbr_material.metallicFactor = 1.0 +# pbr_material.roughnessFactor = 1.0 +# else: +# pbr_material.metallicFactor = 0.0 +# pbr_material.roughnessFactor = 0.9 - if normal_texture is not None: - gltf.model.materials[0].normalTexture = NormalTextureInfo( - index=add_texture(gltf, normal_texture), scale=normal_strength - ) +# if normal_texture is not None: +# gltf.model.materials[0].normalTexture = NormalTextureInfo( +# index=add_texture(gltf, normal_texture), scale=normal_strength +# ) - gltf.model.nodes[0].name = f"mvadapter_node_{task_id}" - gltf.model.meshes[0].name = f"mvadapter_mesh_{task_id}" - gltf.model.materials[0].name = f"mvadapter_material_{task_id}" - gltf.model.asset.generator = "https://github.com/huanngzh/MV-Adapter" - gltf.export(output_path) +# gltf.model.nodes[0].name = f"mvadapter_node_{task_id}" +# gltf.model.meshes[0].name = f"mvadapter_mesh_{task_id}" +# gltf.model.materials[0].name = f"mvadapter_material_{task_id}" +# gltf.model.asset.generator = "https://github.com/huanngzh/MV-Adapter" +# gltf.export(output_path) def replace_mesh_texture_and_save( @@ -508,15 +508,15 @@ def replace_mesh_texture_and_save( texture_format=texture_format, task_id=task_id, ) - elif backend == "gltflib": - replace_mesh_texture_and_save_gltflib( - input_path, - output_path, - texture, - metallic_roughness_texture=metallic_roughness_texture, - normal_texture=normal_texture, - normal_strength=normal_strength, - task_id=task_id, - ) + # elif backend == "gltflib": + # replace_mesh_texture_and_save_gltflib( + # input_path, + # output_path, + # texture, + # metallic_roughness_texture=metallic_roughness_texture, + # normal_texture=normal_texture, + # normal_strength=normal_strength, + # task_id=task_id, + # ) else: raise NotImplementedError diff --git a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py index eddce2b..da88942 100644 --- a/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py +++ b/Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py @@ -281,12 +281,10 @@ def run_pipeline( def create_image_grid(images: List[Image.Image], rows: int = 1) -> Image.Image: - """Создание сетки изображений""" return make_image_grid(images, rows=rows) def pil_to_torch_tensor(image: Image.Image) -> torch.Tensor: - """Конвертация PIL изображения в torch tensor""" import numpy as np array = np.array(image) tensor = torch.from_numpy(array).float() / 255.0 @@ -299,7 +297,6 @@ def pil_to_torch_tensor(image: Image.Image) -> torch.Tensor: def pils_to_torch_tensor(images: List[Image.Image]) -> torch.Tensor: - """Конвертация списка PIL изображений в torch tensor""" tensors = [pil_to_torch_tensor(img).squeeze(0) for img in images] return torch.stack(tensors, dim=0) @@ -453,4 +450,223 @@ def run_t2mv_pipeline( **pipe_kwargs, ).images - return images \ No newline at end of file + return images + + +def prepare_texture_pipeline( + upscaler_ckpt_path: Optional[str] = None, + inpaint_ckpt_path: Optional[str] = None, + device: str = "cuda", + use_mmgp: bool = False, +): + """ + Prepare texture projection pipeline + + Args: + upscaler_ckpt_path: Path to upscaler checkpoint (optional) + inpaint_ckpt_path: Path to inpaint checkpoint (optional) + device: Device for computation + use_mmgp: Use mmgp for memory optimization + """ + from .mvadapter.pipelines.pipeline_texture import TexturePipeline + + if upscaler_ckpt_path and not os.path.exists(upscaler_ckpt_path): + print(f"Warning: Upscaler checkpoint not found: {upscaler_ckpt_path}, using None") + upscaler_ckpt_path = None + + if inpaint_ckpt_path and not os.path.exists(inpaint_ckpt_path): + print(f"Warning: Inpaint checkpoint not found: {inpaint_ckpt_path}, using None") + inpaint_ckpt_path = None + + texture_pipe = TexturePipeline( + upscaler_ckpt_path=upscaler_ckpt_path, + inpaint_ckpt_path=inpaint_ckpt_path, + device=device + ) + + # Apply mmgp optimization if available and requested + if use_mmgp and MMGP_AVAILABLE: + import gc + torch.cuda.empty_cache() + gc.collect() + + all_models = {} + + if hasattr(texture_pipe, 'upscaler') and texture_pipe.upscaler is not None: + all_models["upscaler"] = texture_pipe.upscaler + + if hasattr(texture_pipe, 'inpainter') and texture_pipe.inpainter is not None: + all_models["inpainter"] = texture_pipe.inpainter + + if all_models: + offload.profile(all_models, profile_type.HighRAM_LowVRAM) + print("mmgp profiling applied to texture pipeline successfully.") + else: + print("No models loaded for texture pipeline, skipping mmgp profiling.") + + return texture_pipe + + +def run_texture_pipeline( + texture_pipe, + multiview_images: List[Image.Image], + mesh_path: str, + save_dir: str, + save_name: str = "textured_model", + uv_size: int = 4096, + view_upscale: bool = True, + inpaint_mode: str = "view", + uv_unwarp: bool = True, + preprocess_mesh: bool = False, + move_to_center: bool = False, + front_x: bool = True, + camera_azimuth_deg: List[float] = None, + camera_elevation_deg: List[float] = None, + camera_distance: float = 1.0, + camera_ortho_scale: float = 1.1, + debug_mode: bool = False, + apply_dilate: bool = True, +): + """ + Execute texture projection pipeline + + Args: + texture_pipe: Prepared texture pipeline + multiview_images: List of PIL images (6 views) + mesh_path: Path to 3D mesh file + save_dir: Directory to save results + save_name: Name for saved files + uv_size: UV texture resolution + view_upscale: Enable view upscaling + inpaint_mode: Inpainting mode ("none", "uv", "view") + uv_unwarp: Enable UV unwrapping + preprocess_mesh: Enable mesh preprocessing + move_to_center: Move mesh to center + front_x: Front face along X axis + camera_azimuth_deg: Camera azimuth angles + camera_elevation_deg: Camera elevation angles + camera_distance: Camera distance + camera_ortho_scale: Orthographic scale + debug_mode: Enable debug mode + apply_dilate: Apply dilate to remove texture seams + """ + from .mvadapter.pipelines.pipeline_texture import ModProcessConfig + from .mvadapter.utils import make_image_grid + + if not mesh_path or not os.path.exists(mesh_path): + raise ValueError(f"Mesh file does not exist: {mesh_path}") + + os.makedirs(save_dir, exist_ok=True) + + if camera_azimuth_deg is None: + camera_azimuth_deg = [0, 90, 180, 270, 180, 180] + if camera_elevation_deg is None: + camera_elevation_deg = [0, 0, 0, 0, 89.99, -89.99] + + # Adjust azimuth angles (subtract 90 as in original code) + azimuth_deg_corrected = [x - 90 for x in camera_azimuth_deg] + + grid_image = make_image_grid(multiview_images, rows=1) + temp_grid_path = os.path.join(save_dir, f"{save_name}_temp_grid.png") + grid_image.save(temp_grid_path) + + try: + rgb_process_config = ModProcessConfig( + view_upscale=view_upscale, + inpaint_mode=inpaint_mode + ) + + output = texture_pipe( + mesh_path=mesh_path, + save_dir=save_dir, + save_name=save_name, + move_to_center=move_to_center, + front_x=front_x, + uv_unwarp=uv_unwarp, + preprocess_mesh=preprocess_mesh, + uv_size=uv_size, + rgb_path=temp_grid_path, + rgb_process_config=rgb_process_config, + camera_elevation_deg=camera_elevation_deg, + camera_azimuth_deg=azimuth_deg_corrected, + camera_distance=camera_distance, + camera_ortho_scale=camera_ortho_scale, + debug_mode=debug_mode, + apply_dilate=apply_dilate, + ) + + if os.path.exists(temp_grid_path): + os.remove(temp_grid_path) + + return output + + except Exception as e: + if os.path.exists(temp_grid_path): + os.remove(temp_grid_path) + raise e + + +def download_texture_checkpoints(texture_ckpt_dir, upscaler_ckpt_path, inpaint_ckpt_path): + """Download texture checkpoint files if they don't exist""" + import subprocess + + os.makedirs(texture_ckpt_dir, exist_ok=True) + + checkpoints = [ + { + "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth", + "path": upscaler_ckpt_path, + "name": "RealESRGAN upscaler" + }, + { + "url": "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", + "path": inpaint_ckpt_path, + "name": "Big-LaMa inpainter" + } + ] + + for ckpt in checkpoints: + if not os.path.exists(ckpt["path"]): + print(f"Downloading {ckpt['name']} to {ckpt['path']}...") + + success = False + + # Try wget first + try: + subprocess.run([ + "wget", ckpt["url"], "-O", ckpt["path"] + ], check=True, capture_output=True, text=True) + print(f"Successfully downloaded {ckpt['name']} using wget") + success = True + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + # Try curl if wget failed + if not success: + try: + subprocess.run([ + "curl", "-L", ckpt["url"], "-o", ckpt["path"] + ], check=True, capture_output=True, text=True) + print(f"Successfully downloaded {ckpt['name']} using curl") + success = True + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + # Try Python urllib as fallback + if not success: + try: + import urllib.request + urllib.request.urlretrieve(ckpt["url"], ckpt["path"]) + print(f"Successfully downloaded {ckpt['name']} using urllib") + success = True + except Exception as e: + print(f"Failed to download {ckpt['name']} using urllib: {e}") + + # If all methods failed + if not success: + print(f"Failed to download {ckpt['name']}. Please download manually:") + print(f"wget {ckpt['url']} -O {ckpt['path']}") + print(f"or") + print(f"curl -L {ckpt['url']} -o {ckpt['path']}") + else: + print(f"{ckpt['name']} already exists at {ckpt['path']}") \ No newline at end of file diff --git a/_Example_Workflows/MV-Adapter/MV-Adapter-Texturing.json b/_Example_Workflows/MV-Adapter/MV-Adapter-Texturing.json new file mode 100644 index 0000000..0e82ea2 --- /dev/null +++ b/_Example_Workflows/MV-Adapter/MV-Adapter-Texturing.json @@ -0,0 +1,315 @@ +{ + "id": "4f7b5e42-a56c-480c-8bed-bb1a8a8006cd", + "revision": 0, + "last_node_id": 26, + "last_link_id": 43, + "nodes": [ + { + "id": 22, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 5958.7109375, + -101.48223114013672 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "widget": { + "name": "mesh_file_path" + }, + "link": 42 + } + ], + "outputs": [], + "title": "PBR Mesh", + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 17, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 5346.52099609375, + -95.2625503540039 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "widget": { + "name": "mesh_file_path" + }, + "link": 41 + } + ], + "outputs": [], + "title": "Shaded Mesh", + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 24, + "type": "[Comfy3D] Load MVAdapter Texture Pipeline", + "pos": [ + 5015.30126953125, + -68.61071014404297 + ], + "size": [ + 307.4322204589844, + 130 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "texture_pipeline", + "type": "TEXTURE_PIPE", + "links": [ + 38 + ] + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load MVAdapter Texture Pipeline" + }, + "widgets_values": [ + "RealESRGAN_x2plus.pth", + "big-lama.pt", + false, + true + ] + }, + { + "id": 15, + "type": "LoadImage", + "pos": [ + 4655.7041015625, + 132.22010803222656 + ], + "size": [ + 325.982177734375, + 338.67474365234375 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 39 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "squirel_girl_grid.png", + "image" + ] + }, + { + "id": 26, + "type": "[Comfy3D] MVAdapter Texture Projection", + "pos": [ + 5015.29931640625, + 137.44638061523438 + ], + "size": [ + 298.482421875, + 462 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "texture_pipeline", + "type": "TEXTURE_PIPE", + "link": 38 + }, + { + "name": "grid_image", + "type": "IMAGE", + "link": 39 + }, + { + "name": "mesh_path", + "type": "STRING", + "widget": { + "name": "mesh_path" + }, + "link": 40 + } + ], + "outputs": [ + { + "name": "shaded_model_path", + "type": "STRING", + "links": [ + 41 + ] + }, + { + "name": "pbr_model_path", + "type": "STRING", + "links": [ + 42 + ] + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] MVAdapter Texture Projection" + }, + "widgets_values": [ + "", + "./output", + "textured_model", + 4096, + true, + "view", + true, + false, + false, + true, + true, + true, + "0,90,180,270,180,180", + "0,0,0,0,89.99,-89.99", + 1, + 1.1, + false + ] + }, + { + "id": 4, + "type": "PrimitiveString", + "pos": [ + 4659.12890625, + -19.559751510620117 + ], + "size": [ + 319.9439392089844, + 92.46147918701172 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 40 + ] + } + ], + "title": "Input mesh", + "properties": { + "Node name for S&R": "PrimitiveString" + }, + "widgets_values": [ + "output/StableGenTest/mesh.glb" + ] + } + ], + "links": [ + [ + 38, + 24, + 0, + 26, + 0, + "TEXTURE_PIPE" + ], + [ + 39, + 15, + 0, + 26, + 1, + "IMAGE" + ], + [ + 40, + 4, + 0, + 26, + 2, + "STRING" + ], + [ + 41, + 26, + 0, + 17, + 0, + "STRING" + ], + [ + 42, + 26, + 1, + 22, + 0, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 4.321942375150668, + "offset": [ + -4633.552682414396, + -263.1857272782662 + ] + }, + "frontendVersion": "1.20.4", + "node_versions": { + "ComfyUI-3D-Pack": "626a5ee60d6fbe3d76d69cd89261d445dbca2bc4", + "comfy-core": "0.3.35" + }, + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/MV-Adapter-ig2mv.json b/_Example_Workflows/MV-Adapter/MV-Adapter-ig2mv.json similarity index 64% rename from _Example_Workflows/MV-Adapter-ig2mv.json rename to _Example_Workflows/MV-Adapter/MV-Adapter-ig2mv.json index 606553d..074747f 100644 --- a/_Example_Workflows/MV-Adapter-ig2mv.json +++ b/_Example_Workflows/MV-Adapter/MV-Adapter-ig2mv.json @@ -1,15 +1,15 @@ { "id": "867889cd-eeb4-48b0-853c-21e125a93896", "revision": 0, - "last_node_id": 9, - "last_link_id": 7, + "last_node_id": 28, + "last_link_id": 25, "nodes": [ { "id": 8, "type": "PreviewImage", "pos": [ 3033.075927734375, - -927.3480224609375 + -927.3829956054688 ], "size": [ 370.7532958984375, @@ -32,15 +32,80 @@ "widgets_values": [] }, { - "id": 1, - "type": "[Comfy3D] Load MVAdapter Pipeline", + "id": 26, + "type": "[Comfy3D] Make Image Grid", "pos": [ - 2145.03466796875, - -1135.2872314453125 + 3033.62158203125, + -586.97265625 ], "size": [ - 356.66668701171875, - 250 + 270, + 82 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 22 + } + ], + "outputs": [ + { + "name": "image_grid", + "type": "IMAGE", + "links": [ + 23 + ] + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Make Image Grid" + }, + "widgets_values": [ + 1, + true + ] + }, + { + "id": 27, + "type": "PreviewImage", + "pos": [ + 3032.88818359375, + -449.5286865234375 + ], + "size": [ + 420.0350036621094, + 123.75502014160156 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 23 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 3, + "type": "LoadImage", + "pos": [ + 2153.30078125, + -839.7021484375 + ], + "size": [ + 346.8450622558594, + 326 ], "flags": {}, "order": 0, @@ -48,38 +113,36 @@ "inputs": [], "outputs": [ { - "name": "mvadapter_pipe", - "type": "DIFFUSERS_PIPE", + "name": "IMAGE", + "type": "IMAGE", "links": [ - 4 + 19 ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null } ], "properties": { - "Node name for S&R": "[Comfy3D] Load MVAdapter Pipeline" + "Node name for S&R": "LoadImage" }, "widgets_values": [ - "stabilityai/stable-diffusion-xl-base-1.0", - "madebyollin/sdxl-vae-fp16-fix", - "huanngzh/mv-adapter", - "default", - 6, - true, - true, - "", - "" + "squirrel_girl_1.png", + "image" ] }, { "id": 9, "type": "PrimitiveString", "pos": [ - 2155.025146484375, - -447.5125427246094 + 2155.323974609375, + -452.05352783203125 ], "size": [ - 334.6800231933594, - 66.81999969482422 + 343.5000305175781, + 67.5549545288086 ], "flags": {}, "order": 1, @@ -119,12 +182,12 @@ { "name": "mvadapter_pipe", "type": "DIFFUSERS_PIPE", - "link": 4 + "link": 25 }, { "name": "reference_image", "type": "IMAGE", - "link": 5 + "link": 19 }, { "name": "mesh_path", @@ -140,7 +203,8 @@ "name": "multiview_images", "type": "IMAGE", "links": [ - 6 + 6, + 22 ] } ], @@ -156,22 +220,22 @@ 1, 768, 768, - 92634126158178, + 710491314281528, "randomize", false, 1 ] }, { - "id": 3, - "type": "LoadImage", + "id": 28, + "type": "[Comfy3D] Load MVAdapter IG2MV Pipeline", "pos": [ - 2152.031494140625, - -832.7184448242188 + 2154.555908203125, + -1155.862060546875 ], "size": [ - 346.11004638671875, - 326 + 357.2724609375, + 257.3499755859375 ], "flags": {}, "order": 2, @@ -179,44 +243,30 @@ "inputs": [], "outputs": [ { - "name": "IMAGE", - "type": "IMAGE", + "name": "mvadapter_pipe", + "type": "DIFFUSERS_PIPE", "links": [ - 5 + 25 ] - }, - { - "name": "MASK", - "type": "MASK", - "links": null } ], "properties": { - "Node name for S&R": "LoadImage" + "Node name for S&R": "[Comfy3D] Load MVAdapter IG2MV Pipeline" }, "widgets_values": [ - "squirrel_girl_1.png", - "image" + "stabilityai/stable-diffusion-xl-base-1.0", + "madebyollin/sdxl-vae-fp16-fix", + "huanngzh/mv-adapter", + "default", + 6, + true, + true, + "", + "" ] } ], "links": [ - [ - 4, - 1, - 0, - 7, - 0, - "DIFFUSERS_PIPE" - ], - [ - 5, - 3, - 0, - 7, - 1, - "IMAGE" - ], [ 6, 7, @@ -232,23 +282,55 @@ 7, 2, "STRING" + ], + [ + 19, + 3, + 0, + 7, + 1, + "IMAGE" + ], + [ + 22, + 7, + 0, + 26, + 0, + "IMAGE" + ], + [ + 23, + 26, + 0, + 27, + 0, + "IMAGE" + ], + [ + 25, + 28, + 0, + 7, + 0, + "DIFFUSERS_PIPE" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.9523809523809523, + "scale": 0.9070294784580498, "offset": [ - -1869.4567072737789, - 1193.6790746138436 + -1941.4116551684099, + 1211.3556830868408 ] }, + "frontendVersion": "1.20.4", "node_versions": { "comfy-core": "0.3.35", - "ComfyUI-3D-Pack": "829fc82eeeb63b614510117f0ed9c388c106b5f2" + "ComfyUI-3D-Pack": "626a5ee60d6fbe3d76d69cd89261d445dbca2bc4" }, - "frontendVersion": "1.20.4", "VHS_latentpreview": false, "VHS_latentpreviewrate": 0, "VHS_MetadataImage": true, diff --git a/_Example_Workflows/MV-Adapter-t2mv.json b/_Example_Workflows/MV-Adapter/MV-Adapter-t2mv.json similarity index 100% rename from _Example_Workflows/MV-Adapter-t2mv.json rename to _Example_Workflows/MV-Adapter/MV-Adapter-t2mv.json diff --git a/_Example_Workflows/_Example_Inputs_Files/squirel_girl1_grid.png b/_Example_Workflows/_Example_Inputs_Files/squirel_girl1_grid.png new file mode 100644 index 0000000..6718f8d Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/squirel_girl1_grid.png differ diff --git a/nodes.py b/nodes.py index d9351b9..6e680b1 100644 --- a/nodes.py +++ b/nodes.py @@ -103,13 +103,14 @@ from TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline from TRELLIS.trellis.utils import postprocessing_utils from TripoSG.pipelines.pipeline_triposg import TripoSGPipeline from TripoSG.pipelines.pipeline_triposg_scribble import TripoSGScribblePipeline -from Stable3DGen.trellis.backend_config import get_available_backends, get_available_sparse_backends from Stable3DGen.pipeline_builders import StableGenPipelineBuilder from MV_Adapter.mvadapter_node_utils import ( prepare_pipeline as mvadapter_prepare_pipeline, run_pipeline as mvadapter_run_pipeline, prepare_t2mv_pipeline as mvadapter_prepare_t2mv_pipeline, - run_t2mv_pipeline as mvadapter_run_t2mv_pipeline + run_t2mv_pipeline as mvadapter_run_t2mv_pipeline, + prepare_texture_pipeline as mvadapter_prepare_texture_pipeline, + download_texture_checkpoints, ) from MV_Adapter.mvadapter.utils import make_image_grid @@ -4789,7 +4790,7 @@ class StableGen_StableX_Process_Image: except Exception as e: raise Exception(f"[StableGen_StableX_Process_Image] Image processing failed: {str(e)}") -class Load_MVAdapter_Pipeline: +class Load_MVAdapter_IG2MV_Pipeline: """Loader pipeline for MV-Adapter (Image to Multi-View)""" CATEGORY = "Comfy3D/Algorithm" RETURN_TYPES = ("DIFFUSERS_PIPE",) @@ -5041,3 +5042,176 @@ class MVAdapter_Text_To_MultiView: return_images = pils_to_torch_imgs(images, device=DEVICE_STR) return (return_images,) +class Load_MVAdapter_Texture_Pipeline: + """Load texture projection pipeline for MVAdapter""" + CATEGORY = "Comfy3D/Algorithm" + RETURN_TYPES = ("TEXTURE_PIPE",) + RETURN_NAMES = ("texture_pipeline",) + FUNCTION = "load" + + TEXTURE_CKPT_DIR = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh/MV-Adapter") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "upscaler_ckpt_name": ("STRING", {"default": "RealESRGAN_x2plus.pth"}), + "inpaint_ckpt_name": ("STRING", {"default": "big-lama.pt"}), + "use_mmgp": ("BOOLEAN", {"default": False}), + "auto_download": ("BOOLEAN", {"default": True}), + } + } + + @classmethod + def load(cls, upscaler_ckpt_name, inpaint_ckpt_name, use_mmgp, auto_download): + + # Form full paths from model names + upscaler_ckpt_path = os.path.join(cls.TEXTURE_CKPT_DIR, upscaler_ckpt_name) if upscaler_ckpt_name.strip() else None + inpaint_ckpt_path = os.path.join(cls.TEXTURE_CKPT_DIR, inpaint_ckpt_name) if inpaint_ckpt_name.strip() else None + + # Auto-download checkpoints if enabled + if auto_download: + download_texture_checkpoints(cls.TEXTURE_CKPT_DIR, upscaler_ckpt_path, inpaint_ckpt_path) + + texture_pipe = mvadapter_prepare_texture_pipeline( + upscaler_ckpt_path=upscaler_ckpt_path, + inpaint_ckpt_path=inpaint_ckpt_path, + device=DEVICE_STR, + use_mmgp=use_mmgp, + ) + + print("Texture pipeline loaded successfully") + return (texture_pipe,) + + +class MVAdapter_Texture_Projection: + """Project grid image onto 3D mesh using pre-loaded texture pipeline""" + CATEGORY = "Comfy3D/Algorithm" + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("shaded_model_path", "pbr_model_path") + FUNCTION = "run" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "texture_pipeline": ("TEXTURE_PIPE",), + "grid_image": ("IMAGE",), + "mesh_path": ("STRING", {"default": ""}), + "save_dir": ("STRING", {"default": "./output"}), + "save_name": ("STRING", {"default": "textured_model"}), + "uv_size": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 256}), + "view_upscale": ("BOOLEAN", {"default": True}), + "inpaint_mode": (["none", "uv", "view"], {"default": "view"}), + "uv_unwarp": ("BOOLEAN", {"default": True}), + "preprocess_mesh": ("BOOLEAN", {"default": False}), + "move_to_center": ("BOOLEAN", {"default": False}), + "front_x": ("BOOLEAN", {"default": True}), + "create_pbr_model": ("BOOLEAN", {"default": True}), + "apply_dilate": ("BOOLEAN", {"default": True}), + }, + "optional": { + "camera_azimuth_deg": ("STRING", {"default": "0,90,180,270,180,180"}), + "camera_elevation_deg": ("STRING", {"default": "0,0,0,0,89.99,-89.99"}), + "camera_distance": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "camera_ortho_scale": ("FLOAT", {"default": 1.1, "min": 0.1, "max": 5.0, "step": 0.1}), + "debug_mode": ("BOOLEAN", {"default": False}), + } + } + + def run(self, texture_pipeline, grid_image, mesh_path, save_dir, save_name, + uv_size, view_upscale, inpaint_mode, uv_unwarp, preprocess_mesh, + move_to_center, front_x, create_pbr_model, apply_dilate, + camera_azimuth_deg="0,90,180,270,180,180", + camera_elevation_deg="0,0,0,0,89.99,-89.99", + camera_distance=1.0, camera_ortho_scale=1.1, debug_mode=False): + + if isinstance(grid_image, torch.Tensor): + pil_grid = torch_imgs_to_pils(grid_image)[0] # Get first (and only) image + else: + raise ValueError("grid_image must be torch.Tensor") + + if not mesh_path or not os.path.exists(mesh_path): + raise ValueError(f"Mesh file does not exist: {mesh_path}") + + # Convert to absolute paths + save_dir = os.path.abspath(save_dir) + os.makedirs(save_dir, exist_ok=True) + + # Parse camera angles with fallback to defaults + try: + azimuth_deg = [float(x.strip()) for x in camera_azimuth_deg.split(",")] + elevation_deg = [float(x.strip()) for x in camera_elevation_deg.split(",")] + except: + azimuth_deg = [0, 90, 180, 270, 180, 180] + elevation_deg = [0, 0, 0, 0, 89.99, -89.99] + + # Adjust azimuth angles (subtract 90 as in original code) + azimuth_deg_corrected = [x - 90 for x in azimuth_deg] + + # Save grid image temporarily + temp_grid_path = os.path.join(save_dir, f"{save_name}_temp_grid.png") + pil_grid.save(temp_grid_path) + + try: + from .Gen_3D_Modules.MV_Adapter.mvadapter.pipelines.pipeline_texture import ModProcessConfig + + rgb_process_config = ModProcessConfig( + view_upscale=view_upscale, + inpaint_mode=inpaint_mode + ) + + base_color_process_config = ModProcessConfig( + view_upscale=view_upscale, + inpaint_mode=inpaint_mode + ) + + # Prepare texture pipeline arguments + texture_args = { + "mesh_path": mesh_path, + "save_dir": save_dir, + "save_name": save_name, + "move_to_center": move_to_center, + "front_x": front_x, + "uv_unwarp": uv_unwarp, + "preprocess_mesh": preprocess_mesh, + "uv_size": uv_size, + "rgb_path": temp_grid_path, + "rgb_process_config": rgb_process_config, + "camera_elevation_deg": elevation_deg, + "camera_azimuth_deg": azimuth_deg_corrected, + "camera_distance": camera_distance, + "camera_ortho_scale": camera_ortho_scale, + "debug_mode": debug_mode, + "apply_dilate": apply_dilate, + } + + # Add PBR settings if enabled + if create_pbr_model: + texture_args.update({ + "base_color_path": temp_grid_path, + "base_color_process_config": base_color_process_config, + }) + + output = texture_pipeline(**texture_args) + + # Cleanup temp file + if os.path.exists(temp_grid_path): + os.remove(temp_grid_path) + + # Convert to absolute paths + shaded_path = os.path.abspath(output.shaded_model_save_path) if output.shaded_model_save_path else "" + pbr_path = os.path.abspath(output.pbr_model_save_path) if output.pbr_model_save_path else "" + + print(f"Texturing from grid completed:") + print(f" Shaded model: {shaded_path}") + print(f" PBR model: {pbr_path}") + + return (shaded_path, pbr_path) + + except Exception as e: + # Cleanup temp file on error + if os.path.exists(temp_grid_path): + os.remove(temp_grid_path) + raise e +