Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
129 lines
5.0 KiB
Python
129 lines
5.0 KiB
Python
"""MLX-compatible ViewProcessor for Hunyuan3D texture pipeline.
|
|
|
|
Drop-in replacement for ViewProcessor that works with MeshRenderMLX.
|
|
Uses numpy instead of torch for tensor operations.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
|
|
class ViewProcessorMLX:
|
|
"""Orchestrates multiview rendering and texture baking using MeshRenderMLX."""
|
|
|
|
def __init__(self, config, render):
|
|
self.config = config
|
|
self.render = render
|
|
|
|
def render_normal_multiview(self, camera_elevs, camera_azims, use_abs_coor=True):
|
|
return [
|
|
self.render.render_normal(elev, azim, use_abs_coor=use_abs_coor,
|
|
return_type="pl")
|
|
for elev, azim in zip(camera_elevs, camera_azims)
|
|
]
|
|
|
|
def render_position_multiview(self, camera_elevs, camera_azims):
|
|
return [
|
|
self.render.render_position(elev, azim, return_type="pl")
|
|
for elev, azim in zip(camera_elevs, camera_azims)
|
|
]
|
|
|
|
def bake_view_selection(self, candidate_camera_elevs, candidate_camera_azims,
|
|
candidate_view_weights, max_selected_view_num):
|
|
"""Select optimal camera views for texture baking via greedy set cover."""
|
|
original_resolution = self.render.default_resolution
|
|
self.render.set_default_render_resolution(1024)
|
|
|
|
selected_elevs = []
|
|
selected_azims = []
|
|
selected_weights = []
|
|
|
|
face_areas = self.render.get_face_areas(from_one_index=True)
|
|
total_area = face_areas.sum()
|
|
face_area_ratios = face_areas / total_area
|
|
|
|
n_candidates = len(candidate_camera_elevs)
|
|
self.render.set_boundary_unreliable_scale(2)
|
|
|
|
# Render alpha (face indices) for all candidates
|
|
viewed_tri_idxs = []
|
|
viewed_masks = []
|
|
for elev, azim in zip(candidate_camera_elevs, candidate_camera_azims):
|
|
tri_idx = self.render.render_alpha(elev, azim, return_type="np")
|
|
viewed_tri_idxs.append(set(np.unique(tri_idx.flatten()).tolist()))
|
|
viewed_masks.append(tri_idx[0, :, :, 0] > 0)
|
|
|
|
is_selected = [False] * n_candidates
|
|
total_viewed = set()
|
|
|
|
# Always select first 6 (cardinal directions)
|
|
for idx in range(min(6, n_candidates)):
|
|
selected_elevs.append(candidate_camera_elevs[idx])
|
|
selected_azims.append(candidate_camera_azims[idx])
|
|
selected_weights.append(candidate_view_weights[idx])
|
|
is_selected[idx] = True
|
|
total_viewed.update(viewed_tri_idxs[idx])
|
|
|
|
# Greedy: add views that cover the most new area
|
|
for _ in range(max_selected_view_num - len(selected_weights)):
|
|
max_inc = 0.0
|
|
max_idx = -1
|
|
|
|
for idx in range(n_candidates):
|
|
if is_selected[idx]:
|
|
continue
|
|
new_tris = viewed_tri_idxs[idx] - total_viewed
|
|
if not new_tris:
|
|
continue
|
|
new_area = face_area_ratios[list(new_tris)].sum()
|
|
if new_area > max_inc:
|
|
max_inc = new_area
|
|
max_idx = idx
|
|
|
|
if max_inc > 0.01 and max_idx >= 0:
|
|
is_selected[max_idx] = True
|
|
selected_elevs.append(candidate_camera_elevs[max_idx])
|
|
selected_azims.append(candidate_camera_azims[max_idx])
|
|
selected_weights.append(candidate_view_weights[max_idx])
|
|
total_viewed.update(viewed_tri_idxs[max_idx])
|
|
else:
|
|
break
|
|
|
|
self.render.set_default_render_resolution(original_resolution)
|
|
return selected_elevs, selected_azims, selected_weights
|
|
|
|
def bake_from_multiview(self, views, camera_elevs, camera_azims,
|
|
view_weights):
|
|
"""Back-project multiple views and merge into a single UV texture."""
|
|
textures = []
|
|
cos_maps = []
|
|
|
|
for view, elev, azim, weight in zip(views, camera_elevs, camera_azims,
|
|
view_weights):
|
|
texture, cos_map, _ = self.render.back_project(view, elev, azim)
|
|
cos_map = weight * (cos_map ** self.config.bake_exp)
|
|
textures.append(texture)
|
|
cos_maps.append(cos_map)
|
|
# Incrementally merge (matches original behavior)
|
|
merged, trust = self.render.fast_bake_texture(textures, cos_maps)
|
|
|
|
return merged, trust
|
|
|
|
def texture_inpaint(self, texture, mask, default=None):
|
|
"""Inpaint missing texture regions.
|
|
|
|
Args:
|
|
texture: numpy float32 array or will be converted.
|
|
mask: numpy uint8 mask (255=keep, 0=inpaint).
|
|
default: optional fill value for masked regions.
|
|
|
|
Returns:
|
|
numpy float32 texture.
|
|
"""
|
|
if default is not None:
|
|
bool_mask = mask.astype(bool) if mask.dtype != bool else mask
|
|
texture[~bool_mask] = default
|
|
return texture
|
|
|
|
texture_np = self.render.uv_inpaint(texture, mask)
|
|
return texture_np.astype(np.float32) / 255.0
|