add mv-adapter texturing

This commit is contained in:
Denis Kochetov
2025-06-04 20:16:55 +00:00
parent bd21742093
commit 97975488c6
8 changed files with 1388 additions and 149 deletions
@@ -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,
)
@@ -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
@@ -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
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']}")
@@ -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
}
@@ -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,
Binary file not shown.

After

Width:  |  Height:  |  Size: 5.4 MiB

+177 -3
View File
@@ -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