diff --git a/comfy_extras/mesh3d/uv_unwrap/segment.py b/comfy_extras/mesh3d/uv_unwrap/segment.py index df4ba8ee1..8789b65d9 100644 --- a/comfy_extras/mesh3d/uv_unwrap/segment.py +++ b/comfy_extras/mesh3d/uv_unwrap/segment.py @@ -82,6 +82,7 @@ def segment_charts( w_normal_deviation: float = DEFAULT_W_NORMAL_DEVIATION, w_roundness: float = DEFAULT_W_ROUNDNESS, w_straightness: float = DEFAULT_W_STRAIGHTNESS, + progress_callback=None, ) -> Tensor: """Segment mesh into charts (parallel batch cost-grow). Returns face -> chart_id.""" F = mesh.faces.shape[0] @@ -134,6 +135,7 @@ def segment_charts( thresholds = [t for t in (0.05, 0.1, 0.25) if t < tau_final] + [tau_final] max_inner = max(64, int(F ** 0.5) * 2) outer_iter = 0 + assigned = 0 tq = tqdm(total=F, desc="unwrap: segment (adaptive)", unit="face", leave=False) while True: outer_iter += 1 @@ -147,6 +149,9 @@ def segment_charts( if n_added == 0: break tq.update(n_added) + assigned += n_added + if progress_callback is not None: + progress_callback(assigned, F) unassigned = face_chart < 0 if int(unassigned.sum()) == 0: break @@ -331,6 +336,7 @@ def cluster_charts_pec( max_cost: float = 0.7, max_iters: int = 1024, min_faces: int = 8, + progress_callback=None, ) -> Tensor: """Parallel edge-collapse clustering; returns face_chart [F]. max_cost is the per-merge cutoff (~0.7 rad ~ 40deg); charts under min_faces are then absorbed at a relaxed 2x cutoff.""" @@ -381,6 +387,8 @@ def cluster_charts_pec( n_merge = int(winners.sum().item()) if n_merge == 0: break + if progress_callback is not None: + progress_callback(F - N + n_merge, F) # saturating: charts remaining vs faces win_a = a[winners] win_b = b[winners] diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index adeecb9a0..1923f4c7b 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -2456,6 +2456,8 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance """UV-unwrap a single mesh; returns (vmapping, indices, uvs); vmapping maps each output vertex to an input vertex (seam verts duplicated).""" t_start = time.perf_counter() + # phase-weighted node progress: weld/mesh 2%, segment 33%, extract 5%, param 25%, pack 33% + pbar = comfy.utils.ProgressBar(1000) v_in = positions.to(torch.float32) f_in = indices.to(torch.long).reshape(-1, 3) v_in, f_in, welded_to_orig = _uv_weld_vertices(v_in, f_in, weld_distance) @@ -2471,12 +2473,16 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance logging.warning("[uv_unwrap] mesh face-adjacency < 25% — vertices appear un-welded " "(triangle soup); UV charts will be per-face. Raise weld_distance.") + pbar.update_absolute(20, 1000) + def _seg_progress(done, total): + pbar.update_absolute(20 + (330 * done) // max(total, 1), 1000) if segmenter == "pec": - face_chart = _uv_seg.cluster_charts_pec(mesh, max_cost=1.0) + face_chart = _uv_seg.cluster_charts_pec(mesh, max_cost=1.0, progress_callback=_seg_progress) elif segmenter == "adaptive": - face_chart = _uv_seg.segment_charts(mesh, max_cost=2.0) + face_chart = _uv_seg.segment_charts(mesh, max_cost=2.0, progress_callback=_seg_progress) else: raise ValueError(f"unknown segmenter '{segmenter}'. valid: pec, adaptive") + pbar.update_absolute(350, 1000) n_charts = int(face_chart.max().item()) + 1 if face_chart.numel() else 0 areas_cpu = _uv_mesh.chart_3d_areas(mesh.face_area, face_chart, n_charts).detach().cpu() @@ -2512,8 +2518,7 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance keep = (ff_sorted >= 0) & (face_chart_np[ff_safe] == chart_sorted[:, None]) local_ff_all = np.where(keep, pos_in_chart[ff_safe], -1) - # progress: n_charts units for parameterize + 2*n_charts for pack (prepare + place) - pbar = comfy.utils.ProgressBar(3 * n_charts) + pbar.update_absolute(400, 1000) # parameterize (batched): ortho-project every chart at once, batched stretch metrics # decide acceptance, rejected charts solve ABF/LSCM in dense per-size-bucket batches @@ -2528,7 +2533,8 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance auto = valid_chart & (chart_counts_np <= 5) # tiny charts always keep ortho flip_ok = (n_flip == 0) | (n_flip == chart_counts_np) cand = valid_chart & ~auto & flip_ok & (n_zero == 0) & (rms <= 1.5) & (mx <= 2.0) - pbar.update(int(auto.sum())) + param_done = int(auto.sum()) + pbar.update_absolute(400 + (250 * param_done) // n_charts, 1000) ortho_ok = auto.copy() cand_ids = np.nonzero(cand)[0] @@ -2538,7 +2544,8 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance if not _uv_param._uv_boundary_self_intersects( uv0[v0:v1], local_faces_all[f0:f1], local_ff_all[f0:f1]): ortho_ok[c] = True - pbar.update(1) + param_done += 1 + pbar.update_absolute(400 + (250 * param_done) // n_charts, 1000) lscm_mask = valid_chart & ~ortho_ok batchable = vert_counts <= _uv_param.LSCM_BATCH_MAX_VERTS @@ -2548,7 +2555,8 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance verts_concat, uv0, gl_faces, face_pos, chart_sorted, chart_of_vert, vert_offsets, lscm_ids, n_charts, device=comfy.model_management.get_torch_device()) - pbar.update(int(lscm_ids.size)) + param_done += int(lscm_ids.size) + pbar.update_absolute(400 + (250 * param_done) // n_charts, 1000) uvs_np_list: list = [None] * n_charts uv0_f32 = uv0.astype(np.float32) @@ -2560,7 +2568,8 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance torch.from_numpy(local_faces_all[f0:f1]), torch.from_numpy(local_ff_all[f0:f1]), pin_positions=uv0[v0:v1]) lscm_uv[int(c)] = uvs_t.detach().cpu().numpy().astype(np.float32) - pbar.update(1) + param_done += 1 + pbar.update_absolute(400 + (250 * param_done) // n_charts, 1000) for c in range(n_charts): v0, v1 = vert_offsets[c], vert_offsets[c + 1] if ortho_ok[c]: @@ -2595,13 +2604,13 @@ def _uv_unwrap(positions, indices, segmenter, resolution, padding, weld_distance with tqdm(total=2 * n_charts, desc="unwrap: pack", unit="chart", leave=False) as tq_pack: def _pack_progress(done, total): tq_pack.update(done - tq_pack.n) - pbar.update_absolute(n_charts + done, 3 * n_charts) + pbar.update_absolute(650 + (340 * done) // max(total, 1), 1000) p_x, p_y, p_sw, p_th, p_sc, p_chh, atlas_w, atlas_h = _uv_pack.pack_bitmap_concat( uvs_all_np, vert_offsets, local_faces_all, chart_offsets_np, areas_3d_np, uv_area_np, texels_per_unit=tex_per_unit, padding_texels=padding, progress_callback=_pack_progress) - pbar.update_absolute(3 * n_charts, 3 * n_charts) + pbar.update_absolute(1000, 1000) # assembly: output verts are the per-chart used-vert lists concatenated in chart order, # so vert_offsets doubles as the output vertex cursor