diff --git a/PartPacker/README.md b/PartPacker/README.md new file mode 100644 index 0000000..dcd3017 --- /dev/null +++ b/PartPacker/README.md @@ -0,0 +1,83 @@ +# PartPacker + +![teaser](assets/teaser.gif) + +### [Project Page](https://research.nvidia.com/labs/dir/partpacker/) | [Arxiv](https://arxiv.org/abs/2506.09980) | [Models](https://huggingface.co/nvidia/PartPacker) | [Demo](https://huggingface.co/spaces/nvidia/PartPacker) + + +This is the official implementation of *PartPacker: Efficient Part-level 3D Object Generation via Dual Volume Packing*. + +Our model performs part-level 3D object generation from single-view images. + +### Install + +We rely on `torch` with CUDA installed correctly. + +```bash +pip install -r requirements.txt + +# if you prefer fixed version of dependencies: +pip install -r requirements.lock.txt + +# by default we use torch's built-in attention, if you want to explicitly use flash-attn: +pip install flash-attn --no-build-isolation + +# if you want to run data processing and vae inference, please install meshiki: +pip install meshiki +``` + +### Pretrained models + +Download the pretrained models from huggingface, and put them in the `pretrained` folder. + +```bash +mkdir pretrained +cd pretrained +wget https://huggingface.co/nvidia/PartPacker/resolve/main/vae.pt +wget https://huggingface.co/nvidia/PartPacker/resolve/main/flow.pt +``` + +### Inference + +For inference, it takes ~16GB GPU memory (assuming float16). + +```bash +# vae reconstruction of meshes +PYTHONPATH=. python vae/scripts/infer.py --ckpt_path pretrained/vae.pt --input assets/meshes/ --output_dir output/ + +# flow 3D generation from images +PYTHONPATH=. python flow/scripts/infer.py --ckpt_path pretrained/flow.pt --input assets/images/ --output_dir output/ + +# open local gradio app +python app.py +``` + + +### Data Processing + +We provide a *Dual Volume Packing* implementation to process raw glb meshes into two separate meshes as proposed in the paper. + +```bash +cd data +python bipartite_contraction.py ./example_mesh.glb +# the two separate meshes will be saved in ./output +``` + +### Acknowledgements + +This work is built on many amazing research works and open-source projects, thanks a lot to all the authors for sharing! + +* [Dora](https://github.com/Seed3D/Dora) +* [Hunyuan3D-2](https://github.com/Tencent/Hunyuan3D-2) +* [Trellis](https://github.com/microsoft/TRELLIS) + +## Citation + +``` +@article{tang2024partpacker, + title={Efficient Part-level 3D Object Generation via Dual Volume Packing}, + author={Tang, Jiaxiang and Lu, Ruijie and Li, Zhaoshuo and Hao, Zekun and Li, Xuan and Wei, Fangyin and Song, Shuran and Zeng, Gang and Liu, Ming-Yu and Lin, Tsung-Yi}, + journal={arXiv preprint arXiv:2506.09980}, + year={2025} +} +``` diff --git a/PartPacker/app.py b/PartPacker/app.py new file mode 100644 index 0000000..6ea1742 --- /dev/null +++ b/PartPacker/app.py @@ -0,0 +1,106 @@ +import os +from datetime import datetime + +import cv2 +import kiui +import numpy as np +import rembg +import torch +import trimesh + +from .flow.utils import get_random_color, recenter_foreground +from .vae.utils import postprocess_mesh + + +MAX_SEED = np.iinfo(np.int32).max + + +# get random seed +def get_random_seed(randomize_seed, seed): + if randomize_seed: + seed = np.random.randint(0, MAX_SEED) + return seed + + +# # process image +# @spaces.GPU(duration=10) +def process_image(image,bg_remover,mask_): #cv2 RGB + + # bg removal if there is no alpha channel + if mask_ is None: + image = rembg.remove(image, session=bg_remover) # [H, W, 4] + mask = image[..., -1] > 0 + image = recenter_foreground(image, mask, border_ratio=0.1) + else: + image=mask_ + image = cv2.resize(image, (518, 518), interpolation=cv2.INTER_AREA) + return image + + +# # process generation +# @spaces.GPU(duration=90) +def process_3d(model,bg_remover,input_image,TRIMESH_GLB_EXPORT,mask,output_dir, num_steps=50, cfg_scale=7, grid_res=384, seed=42, simplify_mesh=False, target_num_faces=100000): + + # seed + kiui.seed_everything(seed) + input_image=process_image(input_image,bg_remover,mask) + # output path + # os.makedirs("output", exist_ok=True) + output_glb_path=os.path.join(output_dir, f"partpacker_{datetime.now().strftime('%Y%m%d_%H%M%S')}.glb") + + # input image (assume processed to RGBA uint8) + img_cv2=cv2.cvtColor(input_image,cv2.COLOR_RGB2BGR) + cv2.imwrite("output_image.png", img_cv2) + + image = input_image.astype(np.float32) / 255.0 + if mask is None: + image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4]) # white background + + image_tensor = torch.from_numpy(image).permute(2, 0, 1).contiguous().unsqueeze(0).float().cuda() + + data = {"cond_images": image_tensor} + + with torch.inference_mode(): + results = model(data, num_steps=num_steps, cfg_scale=cfg_scale) + + latent = results["latent"] + + # query mesh + + data_part0 = {"latent": latent[:, : model.config.latent_size, :]} + data_part1 = {"latent": latent[:, model.config.latent_size :, :]} + + with torch.inference_mode(): + results_part0 = model.vae(data_part0, resolution=grid_res) + results_part1 = model.vae(data_part1, resolution=grid_res) + + if not simplify_mesh: + target_num_faces = -1 + + vertices, faces = results_part0["meshes"][0] + mesh_part0 = trimesh.Trimesh(vertices, faces) + mesh_part0.vertices = mesh_part0.vertices @ TRIMESH_GLB_EXPORT.T + mesh_part0 = postprocess_mesh(mesh_part0, target_num_faces) + parts = mesh_part0.split(only_watertight=False) + + vertices, faces = results_part1["meshes"][0] + mesh_part1 = trimesh.Trimesh(vertices, faces) + mesh_part1.vertices = mesh_part1.vertices @ TRIMESH_GLB_EXPORT.T + mesh_part1 = postprocess_mesh(mesh_part1, target_num_faces) + parts.extend(mesh_part1.split(only_watertight=False)) + + # some parts only have 1 face, seems a problem of trimesh.split. + parts = [part for part in parts if len(part.faces) > 10] + + # split connected components and assign different colors + for j, part in enumerate(parts): + # each component uses a random color + part.visual.vertex_colors = get_random_color(j, use_float=True) + + mesh = trimesh.Scene(parts) + # export the whole mesh + mesh.export(output_glb_path) + + return output_glb_path + + diff --git a/PartPacker/assets/images/barrel.png b/PartPacker/assets/images/barrel.png new file mode 100644 index 0000000..5913b29 Binary files /dev/null and b/PartPacker/assets/images/barrel.png differ diff --git a/PartPacker/assets/images/cactus.png b/PartPacker/assets/images/cactus.png new file mode 100644 index 0000000..d6e18c2 Binary files /dev/null and b/PartPacker/assets/images/cactus.png differ diff --git a/PartPacker/assets/images/cyan_car.png b/PartPacker/assets/images/cyan_car.png new file mode 100644 index 0000000..a3332b8 Binary files /dev/null and b/PartPacker/assets/images/cyan_car.png differ diff --git a/PartPacker/assets/images/pickup.png b/PartPacker/assets/images/pickup.png new file mode 100644 index 0000000..dfff502 Binary files /dev/null and b/PartPacker/assets/images/pickup.png differ diff --git a/PartPacker/assets/images/rabbit.png b/PartPacker/assets/images/rabbit.png new file mode 100644 index 0000000..e76806f Binary files /dev/null and b/PartPacker/assets/images/rabbit.png differ diff --git a/PartPacker/assets/images/robot.png b/PartPacker/assets/images/robot.png new file mode 100644 index 0000000..45d4390 Binary files /dev/null and b/PartPacker/assets/images/robot.png differ diff --git a/PartPacker/assets/images/swivelchair.png b/PartPacker/assets/images/swivelchair.png new file mode 100644 index 0000000..75615c1 Binary files /dev/null and b/PartPacker/assets/images/swivelchair.png differ diff --git a/PartPacker/assets/images/teapot.png b/PartPacker/assets/images/teapot.png new file mode 100644 index 0000000..127d167 Binary files /dev/null and b/PartPacker/assets/images/teapot.png differ diff --git a/PartPacker/assets/images/warhammer.png b/PartPacker/assets/images/warhammer.png new file mode 100644 index 0000000..a27e1a0 Binary files /dev/null and b/PartPacker/assets/images/warhammer.png differ diff --git a/PartPacker/assets/meshes/balloon_whisk.ply b/PartPacker/assets/meshes/balloon_whisk.ply new file mode 100644 index 0000000..1fd3a67 Binary files /dev/null and b/PartPacker/assets/meshes/balloon_whisk.ply differ diff --git a/PartPacker/assets/teaser.gif b/PartPacker/assets/teaser.gif new file mode 100644 index 0000000..481a4b6 Binary files /dev/null and b/PartPacker/assets/teaser.gif differ diff --git a/PartPacker/data/bipartite_contraction.py b/PartPacker/data/bipartite_contraction.py new file mode 100644 index 0000000..6b05ccd --- /dev/null +++ b/PartPacker/data/bipartite_contraction.py @@ -0,0 +1,644 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import sys + +sys.path.append(".") + +import argparse +import glob +import os + +import kiui +import numpy as np +import tqdm +import trimesh +from meshiki import Mesh + +parser = argparse.ArgumentParser() +parser.add_argument("test_path", type=str, help="path to the mesh file or folder") +parser.add_argument("--verbose", action="store_true", help="print verbose output") +parser.add_argument("--force_cc", action="store_true", help="force to use connected components and ignore glb groups") +parser.add_argument("--no_smart_group", action="store_true", help="do not perform smart grouping") +parser.add_argument("--no_stitch", action="store_true", help="do not stitch open boundaries") +parser.add_argument("--no_merge_odd_loops", action="store_true", help="do not merge odd loops") +parser.add_argument("--no_dilate", action="store_true", help="do not dilate the mesh") +parser.add_argument("--dilate_size", type=float, default=2 / 512, help="dilate size") +parser.add_argument("--workspace", type=str, default="output", help="path to the output folder") +opt = parser.parse_args() + + +class NamedDisjointSet: + def __init__(self, names): + # names: list of str + self.parent = {name: name for name in names} + + def find(self, x): + if self.parent[x] != x: + self.parent[x] = self.find(self.parent[x]) + return self.parent[x] + + def merge(self, x, y): + self.parent[self.find(x)] = self.find(y) + + def get_groups(self): + groups = {} + for name in self.parent: + root = self.find(name) + if root not in groups: + groups[root] = [] + groups[root].append(name) + return groups + + +def is_single_layer_plane(mesh: trimesh.Trimesh, coplane_thresh: float = 1): + # check if a mesh is just a single-layer plane + if mesh.is_watertight: + return False + face_normals = mesh.face_normals + diff = np.linalg.norm(np.abs(face_normals) - np.abs(face_normals[0])) + return diff.max() < coplane_thresh + + +def calc_intersection_union(bounds1, bounds2): + # bounds: [2, 3] + bmin1, bmax1 = bounds1[0], bounds1[1] + bmin2, bmax2 = bounds2[0], bounds2[1] + # intersection + bmin = np.maximum(bmin1, bmin2) + bmax = np.minimum(bmax1, bmax2) + if np.any(bmin >= bmax): + intersection = 0 + else: + intersection = np.prod(bmax - bmin) + # union + vol1 = np.prod(bmax1 - bmin1) + vol2 = np.prod(bmax2 - bmin2) + union = vol1 + vol2 - intersection + return intersection, union + + +def is_coplanar_and_convex(vertices, coplanar_thresh: float = 1): + # vertices: [N, 3], assume ordered to form a coplanar polygon + # note the last vertex is the same as the first vertex, i.e. ABCA + # we are actually not requiring perfect coplanarity, the thresh of 1 means 60 degree tolerance... + + # Need at least 3 vertices to form a polygon + if len(vertices) < 3: + return False + + # Convert to numpy array if not already + points = np.array(vertices) + n_points = len(points) + + # Normal of the first triangle + v1 = points[1] - points[0] + v2 = points[2] - points[0] + normal = np.cross(v1, v2) + normal = normal / (np.linalg.norm(normal) + 1e-12) + + # Test if coplanar using the normal of fan-cut triangles + for i in range(2, n_points - 2): + v1 = points[i] - points[0] + v2 = points[i + 1] - points[0] + normal_cur = np.cross(v1, v2) + normal_cur = normal_cur / (np.linalg.norm(normal_cur) + 1e-12) + diff = np.linalg.norm(np.abs(normal_cur) - np.abs(normal)) + if diff > coplanar_thresh: + # print(f'not coplanar: {normal_cur} != {normal} (diff = {diff:.4f}) at {i}-{i+1}') + return False # not coplanar + + # Find basis vectors for the 2D plane + # First basis vector can be the normalized vector from points[0] to points[1] + basis1 = v1 / (np.linalg.norm(v1) + 1e-12) + # Second basis vector is perpendicular to both normal and basis1 + basis2 = np.cross(normal, basis1) + basis2 = basis2 / (np.linalg.norm(basis2) + 1e-12) + + # Project all points onto the 2D plane + points_2d = np.zeros((n_points, 2)) + for i in range(n_points): + v = points[i] - points[0] + points_2d[i, 0] = np.dot(v, basis1) + points_2d[i, 1] = np.dot(v, basis2) + + # Check if polygon is convex by using the cross product + # For a convex polygon, all cross products should have the same sign + sign = 0 + for i in range(n_points - 1): + j = (i + 1) % n_points + k = (i + 2) % n_points + + # Vectors from point i to j and j to k + v1 = points_2d[j] - points_2d[i] + v2 = points_2d[k] - points_2d[j] + + # 2D cross product + cross_product = v1[0] * v2[1] - v1[1] * v2[0] + + # Check for consistent sign of cross product + if abs(cross_product) > 1e-2: # Skip collinear points + current_sign = np.sign(cross_product) + if sign == 0: + sign = current_sign + elif sign != current_sign: + # print(f'not convex: {i}-{j} and {j}-{k}, cross_product = {cross_product:.4f}') + return False # not convex + + return True + + +def stitch_nonwatertight_mesh(mesh: trimesh.Trimesh, eps: float = 1e-2): + # mesh will be inplace modified + # return a flag denoting if there are still open boundaries unfixed + + # manager = trimesh.collision.CollisionManager() + # manager.add_object('main', mesh) + + # watertight mesh doesn't need to be stitched + if mesh.is_watertight: + return False + + # planar mesh cannot be stitched + if is_single_layer_plane(mesh, eps): + return True + + # the following is modified from trimesh.repair.stitch + # fan_faces = trimesh.repair.stitch(mesh) + + nonwatertight = False + + from trimesh.path.exchange.misc import faces_to_path + + faces = np.arange(len(mesh.faces)) + + # get a sequence of vertex indices representing the + # boundary of the specified faces + # will be referencing the same indexes of `mesh.vertices` + boundaries = [ + e.points for e in faces_to_path(mesh, faces)["entities"] if len(e.points) > 3 and e.points[0] == e.points[-1] + ] + + # get properties to avoid querying in loop + vertices = mesh.vertices + normals = mesh.face_normals + + # find which faces are associated with an edge + edges_face = mesh.edges_face + tree_edge = mesh.edges_sorted_tree + + # MODIFIED: if any two boundary edges share close vertices, we discard both since they may connect + mask = np.ones(len(boundaries), dtype=bool) + for i in range(len(boundaries)): + for j in range(i + 1, len(boundaries)): + verts_i = vertices[boundaries[i]] # [N, 3] + verts_j = vertices[boundaries[j]] # [M, 3] + # check pair-wise distance + dists = np.linalg.norm(verts_i[:, None, :] - verts_j[None, :, :], axis=-1) # [N, M] + num_close = np.sum(dists < 1e-6) + if num_close >= 4 or (num_close / verts_i.shape[0] >= 0.5) or (num_close / verts_j.shape[0] >= 0.5): + mask[i] = False + mask[j] = False + # print(f'discarding boundary {i} and {j} because of close vertices') + boundaries = [boundaries[i] for i in range(len(boundaries)) if mask[i]] + + # MODIFIED: we only keep coplanar & convex fans + fans = [] + for vert_indices in boundaries: + + # the fan should be coplanar and convex + verts = vertices[vert_indices] + if not is_coplanar_and_convex(verts): + nonwatertight = True + continue + + fan = np.column_stack( + (np.ones(len(vert_indices) - 3, dtype=int) * vert_indices[0], vert_indices[1:-2], vert_indices[2:-1]) + ) # [N, 3] + + fans.append(fan) + + # now we do a normal check against an adjacent face + # to see if each region needs to be flipped + for i, t in zip(range(len(fans)), fans): + # get the edges from the original mesh + # for the first `n` new triangles + e = t[:10, 1:].copy() + e.sort(axis=1) + + # find which indexes of `mesh.edges` these + # new edges correspond with by finding edges + # that exactly correspond with the tree + query = tree_edge.query_ball_point(e, r=1e-10) + if len(query) == 0: + continue + # stack all the indices that exist + edge_index = np.concatenate(query) + + # get the normals from the original mesh + original = normals[edges_face[edge_index]] + + # calculate the normals for a few new faces + check, valid = trimesh.triangles.normals(vertices[t[:3]]) + if not valid.any(): + continue + # take the first valid normal from our new faces + check = check[0] + + # if our new faces are reversed from the original + # Adjacent face flip them along their axis + sign = np.dot(original, check) + if sign.mean() < 0: + fans[i] = np.fliplr(t) + + if len(fans) > 0: + fans = np.vstack(fans) + mesh.faces = np.concatenate([mesh.faces, fans]) + + return nonwatertight + + +def smart_grouping(meshes: dict): + # meshes: {name: trimesh.Trimesh, ...} + + # use collision manager to find all colliding pairs + manager = trimesh.collision.CollisionManager() + for name, mesh in meshes.items(): + manager.add_object(name, mesh) + + is_collide, collide_pairs = manager.in_collision_internal(return_names=True) + # print(f'[INFO] num_collide = {len(collide_pairs)}, {collide_pairs}') + + if not is_collide: + return meshes + + # pre-calculate some stat for each mesh + name_to_stat = {} + total_volume = 0 + max_extent = 0 + num_submeshes = 0 + num_meshes = len(meshes) + for name, mesh in meshes.items(): + name_to_stat[name] = {} + submeshes = mesh.split() + name_to_stat[name]["volume"] = [] + name_to_stat[name]["extent"] = [] + num_submeshes += len(submeshes) + for submesh in submeshes: + if submesh.is_watertight: + name_to_stat[name]["volume"].append(submesh.volume) + total_volume += name_to_stat[name]["volume"][-1] + name_to_stat[name]["extent"].append(np.max(submesh.extents)) + max_extent = max(max_extent, name_to_stat[name]["extent"][-1]) + name_to_stat[name]["volume"] = ( + np.mean(name_to_stat[name]["volume"]) if len(name_to_stat[name]["volume"]) > 0 else np.inf + ) + name_to_stat[name]["extent"] = ( + np.max(name_to_stat[name]["extent"]) if len(name_to_stat[name]["extent"]) > 0 else np.inf + ) + + # use a disjoint set to record grouping + ds = NamedDisjointSet(list(meshes.keys())) + + # decide the merging thresh adaptively (based on the number of meshes, average volume and extent) + # very empirical... + if num_meshes <= 16: + tol_volume = 0.05 * total_volume + tol_extent = 0.05 * max_extent + else: + tol_volume = 0.1 * total_volume + tol_extent = 0.1 * max_extent + + # loop each pair, determine if they should be grouped together + for name1, name2 in collide_pairs: + + if ds.find(name1) == ds.find(name2): + continue + + # single-layer plane should be merged + if is_single_layer_plane(meshes[name1]) or is_single_layer_plane(meshes[name2]): + # print(f'[INFO] merge {name1} and {name2} because of single-layer plane') + ds.merge(name1, name2) + continue + + # too small component + if (name_to_stat[name1]["volume"] < tol_volume and name_to_stat[name1]["extent"] < tol_extent) or ( + name_to_stat[name2]["volume"] < tol_volume and name_to_stat[name2]["extent"] < tol_extent + ): + # print(f'[INFO] merge {name1} and {name2} because of small volume') + ds.merge(name1, name2) + continue + + # overlaps a lot should be merged (just use bounding box IoU) + bounds1 = meshes[name1].bounds # [2, 3] + bounds2 = meshes[name2].bounds # [2, 3] + vol_intersect, vol_union = calc_intersection_union(bounds1, bounds2) + vol_iou = vol_intersect / vol_union + if vol_iou > 0.5: + # print(f'[INFO] merge {name1} and {name2} because of large IoU') + ds.merge(name1, name2) + continue + + # merge groups + groups = ds.get_groups() + for group in groups.values(): + if len(group) <= 1: + continue + # print(f'[INFO] merge group: {group}') + new_name = "_".join(group) + new_mesh = trimesh.util.concatenate(list(meshes[name] for name in group)) + + # merge close vertices, and clean up + new_mesh.merge_vertices(merge_tex=True, merge_norm=True) + new_mesh.update_faces(new_mesh.unique_faces() & new_mesh.nondegenerate_faces()) + new_mesh.fix_normals() + + meshes[new_name] = new_mesh + for name in group: + del meshes[name] + + print(f"[INFO] after grouping, num_meshes = {len(meshes)}") + return meshes + + +def merge_odd_loops(meshes: dict, graph: dict, penetration_depths: dict): + # meshes: {name: trimesh.Trimesh, ...} + # graph: {name: [neighbor, ...], ...} + # penetration_depths: {(name1, name2): depth, ...} + + # find out odd loops in the graph, and try to merge vertex pair of largest penetration depth to make them even loops + all_loops = [] + visited = set() + + def dfs(node, parent, start_node, path): + visited.add(node) + path.append(node) + for neighbor in graph[node]: + if neighbor == parent: + continue + if neighbor == start_node and len(path) >= 3: + all_loops.append(path.copy()) + elif neighbor not in visited: + dfs(neighbor, node, start_node, path) + path.pop() + visited.remove(node) + + for node in graph: + dfs(node, None, node, []) + visited.add(node) + + unique_loops = [] + loop_keys = set() + for loop in all_loops: + loop_key = frozenset(loop) + if len(loop) > 2 and loop_key not in loop_keys: + loop_keys.add(loop_key) + unique_loops.append(loop) + # print(f'[INFO] find loop: {loop}') + + # convert loops to list of edges + loop_edges = [] + for loop in unique_loops: + edges = set() + for i in range(len(loop)): + j = (i + 1) % len(loop) + edges.add(tuple(sorted([loop[i], loop[j]]))) + loop_edges.append(edges) + + # merge vertex pair of largest penetration depth to make them even loops + ds = NamedDisjointSet(list(meshes.keys())) + + while True: # have to loop multiple times to make sure there is no odd loop... + + for i in range(len(loop_edges)): + edges = loop_edges[i] + if len(edges) % 2 == 0: # even + continue + + # find the largest penetration depth + max_penetration_depth = 0 + max_edge = None + for edge in edges: + if penetration_depths[edge] > max_penetration_depth: + max_penetration_depth = penetration_depths[edge] + max_edge = edge + + # remove this edge from other loops if it exists (this will affect other loops, even make already even loops become odd) + for j in range(len(loop_edges)): + if max_edge in loop_edges[j]: + loop_edges[j].remove(max_edge) + + # merge the vertex pair of largest penetration depth + # print(f'[INFO] merge {max_edge[0]} and {max_edge[1]}') + ds.merge(max_edge[0], max_edge[1]) + + has_odd_loop = False + for i in range(len(loop_edges)): + edges = loop_edges[i] + if len(edges) % 2 == 1: + has_odd_loop = True + break + + if not has_odd_loop: + break + + # merge groups + graph_new = graph.copy() + groups = ds.get_groups() + for group in groups.values(): + if len(group) <= 1: + continue + # print(f'[INFO] merge group: {group}') + new_name = "_".join(group) + new_mesh = trimesh.util.concatenate(list(meshes[name] for name in group)) + meshes[new_name] = new_mesh + graph_new[new_name] = set() + for name in group: + del meshes[name] + for neighbor in graph[name]: + if neighbor not in group: + graph_new[new_name].add(neighbor) + graph_new[neighbor].remove(name) + graph_new[neighbor].add(new_name) + del graph_new[name] + + print(f"[INFO] after merging odd loops, num_meshes = {len(meshes)}") + return meshes, graph_new + + +def run(path): + print(f"[INFO] processing {path}") + + mesh = trimesh.load(path) + + if not opt.force_cc and isinstance(mesh, trimesh.Scene) and len(mesh.geometry) > 1: + print(f"[INFO] scene: {len(mesh.geometry)} meshes") + scene = mesh + + else: + if isinstance(mesh, trimesh.Scene): + mesh = mesh.to_mesh() + # use meshiki backend + mesh = Mesh(mesh.vertices, mesh.faces, verbose=opt.verbose, clean=False) + mesh.smart_group_components() + scene = mesh.export_components_as_trimesh_scene() + + # use trimesh + # meshes = mesh.split(only_watertight=False) + # # print(meshes) + # scene = trimesh.Scene() + # for mesh in meshes: + # scene.add_geometry(mesh) + print(f"[INFO] mesh: {len(scene.geometry)} components") + + ### box normalize into [-1, 1] + bounds = scene.bounds # [2, 3] + center = scene.centroid # [3] + # print(f'[INFO] center = {center}, bounds = {bounds}') + scale = 0.95 * 1 / np.max(bounds[1] - bounds[0]) + transform_normalize = np.eye(4) + transform_normalize[:3, 3] = -center + transform_normalize[:3, :3] = np.diag(np.array([scale, scale, scale])) + # print(transform_normalize) + scene.apply_transform(transform_normalize) + + ### apply transform to vertices + meshes = {} + scene_graph = scene.graph.to_flattened() + for k, v in scene_graph.items(): + name = v["geometry"] + if name in scene.geometry and isinstance(scene.geometry[name], trimesh.Trimesh): + transform = v["transform"] + mesh: trimesh.Trimesh = scene.geometry[name].apply_transform(transform) + # drop all textures since we only need geom for parted data + mesh.visual = trimesh.visual.ColorVisuals() + # clean up + mesh.merge_vertices(merge_tex=True, merge_norm=True) + mesh.update_faces(mesh.unique_faces() & mesh.nondegenerate_faces()) + mesh.fix_normals() + meshes[name] = mesh + + ### smart grouping to avoid too many single-layer surface or too small objects + if not opt.no_smart_group: + meshes = smart_grouping(meshes) + + ### stitch open boundaries to make each mesh watertight + if not opt.no_stitch: + for name, mesh in meshes.items(): + stitch_nonwatertight_mesh(mesh) + + ### coloring + # build an undirected collision graph + manager = trimesh.collision.CollisionManager() + for name, mesh in meshes.items(): + # scale up the mesh a little bit to take count of the collision margin + mesh_dilated = mesh.copy() + if not opt.no_dilate: + center = mesh_dilated.centroid + vertices = mesh_dilated.vertices - center + max_radius = np.max(np.linalg.norm(vertices, axis=-1)) + scale = (max_radius + opt.dilate_size) / max_radius + # print(f'[INFO] dilate {name} by {scale}') + mesh_dilated.vertices = vertices * scale + center + manager.add_object(name, mesh_dilated) + + is_collide, collide_pairs, collide_data = manager.in_collision_internal(return_names=True, return_data=True) + + graph = {name: set() for name in meshes.keys()} + penetration_depths = {} + + for data in collide_data: + name1, name2 = list(data.names) + graph[name1].add(name2) + graph[name2].add(name1) + name_key = tuple(sorted([name1, name2])) + penetration_depth = data.depth + penetration_depths[name_key] = penetration_depth + + # merge odd loops + if not opt.no_merge_odd_loops: + # if the graph is too complex, we will skip since it takes forever + num_edges = sum(len(edges) for edges in graph.values()) + if num_edges > 100: + print(f"[WARN] skip {path} because of too many edges: {num_edges}") + else: + meshes, graph = merge_odd_loops(meshes, graph, penetration_depths) + + if opt.verbose: + print(graph) + + # sort objects by distance to center + name_to_centers = {} + for name, mesh in meshes.items(): + vmin = np.min(mesh.vertices, axis=0) + vmax = np.max(mesh.vertices, axis=0) + name_to_centers[name] = (vmin + vmax) / 2 + + name_with_dist = [] # [(name, dist), ...] + for name, mesh_center in name_to_centers.items(): + dist = np.linalg.norm(mesh_center) + name_with_dist.append((name, dist)) + name_with_dist.sort(key=lambda x: x[1]) + + # we will start graph coloring from center to border + name_to_color = {} + queue = [] + initial_color = 0 + for name, dist in name_with_dist: + if name not in name_to_color: + name_to_color[name] = initial_color + initial_color = 1 - initial_color + queue.append(name) + while len(queue) > 0: + name = queue.pop(0) + for neighbor in graph[name]: + if neighbor not in name_to_color: + name_to_color[neighbor] = 1 - name_to_color[name] + if opt.verbose: + print(f"[INFO] color {neighbor} with {name_to_color[neighbor]}") + queue.append(neighbor) + else: + if name_to_color[neighbor] == name_to_color[name]: + print(f"[WARN] {name} and {neighbor} have the same color!") + + # get the two parts + mesh_color0 = [] + mesh_color1 = [] + for name, color in name_to_color.items(): + if color == 0: + mesh_color0.append(meshes[name]) + else: + mesh_color1.append(meshes[name]) + + ### convert to a single mesh and export as glb + mesh_color0 = trimesh.util.concatenate(mesh_color0) + mesh_color1 = trimesh.util.concatenate(mesh_color1) + name = os.path.splitext(os.path.basename(path))[0] + + # export separately + mesh_color0.export(f"{opt.workspace}/{name}_color0.obj") + mesh_color1.export(f"{opt.workspace}/{name}_color1.obj") + + # export together (offsetted) + mesh_color1.vertices += [0, 0, 1] + mesh_all = trimesh.util.concatenate([mesh_color0, mesh_color1]) + mesh_all.export(f"{opt.workspace}/{name}.obj") + + +os.makedirs(opt.workspace, exist_ok=True) + +if os.path.isdir(opt.test_path): + file_paths = glob.glob(os.path.join(opt.test_path, "*")) + for path in tqdm.tqdm(file_paths): + run(path) +else: + run(opt.test_path) diff --git a/PartPacker/docker/Dockerfile b/PartPacker/docker/Dockerfile new file mode 100644 index 0000000..e194596 --- /dev/null +++ b/PartPacker/docker/Dockerfile @@ -0,0 +1,128 @@ +# get the development image from nvidia cuda 12.4 (using devel for full CUDA toolkit) +FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 + +LABEL name="partpacker" maintainer="partpacker" + +# create workspace folder and set it as working directory +RUN mkdir -p /workspace +WORKDIR /workspace + +# update package lists and install essential packages +RUN apt-get update && apt-get install -y \ + build-essential \ + git \ + wget \ + vim \ + libegl1-mesa-dev \ + libglib2.0-0 \ + unzip \ + git-lfs \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Install additional graphics and rendering dependencies +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + pkg-config \ + libglvnd0 \ + libgl1 \ + libglx0 \ + libegl1 \ + libgles2 \ + libglvnd-dev \ + libgl1-mesa-dev \ + libegl1-mesa-dev \ + libgles2-mesa-dev \ + cmake \ + mesa-utils-extra \ + libxrender1 \ + libxi6 \ + libgconf-2-4 \ + libxkbcommon-x11-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +# Set environment variables +ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONUNBUFFERED=1 +ENV LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH +ENV PYOPENGL_PLATFORM=egl + +# Set CUDA environment variables +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=${CUDA_HOME}/bin:${PATH} +ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} +ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;8.9;9.0" + +# install conda +RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \ + chmod +x Miniconda3-latest-Linux-x86_64.sh && \ + ./Miniconda3-latest-Linux-x86_64.sh -b -p /workspace/miniconda3 && \ + rm Miniconda3-latest-Linux-x86_64.sh + +# update PATH environment variable +ENV PATH="/workspace/miniconda3/bin:${PATH}" + +# initialize conda +RUN conda init bash + +# create and activate conda environment +RUN conda create -n partpacker python=3.10 && echo "source activate partpacker" > ~/.bashrc +ENV PATH="/workspace/miniconda3/envs/partpacker/bin:${PATH}" + +# Set conda to always auto-approve +RUN conda config --set always_yes true + +# Install essential conda packages +RUN conda install Ninja +RUN conda install cuda -c nvidia/label/cuda-12.4.1 -y + +# Update libstdcxx-ng to fix compatibility issues +RUN conda install -c conda-forge libstdcxx-ng -y + +# Install PyTorch with CUDA support +RUN pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124 + +# Clone PartPacker repository +RUN git clone https://github.com/NVlabs/PartPacker.git + +# Set working directory to the cloned repository +WORKDIR /workspace/PartPacker + +# Clean up requirements.txt to remove invalid pip options +RUN sed -i 's/ --no-build-isolation//g' requirements.txt && \ + sed -i 's/--no-build-isolation//g' requirements.txt + +# Install Python dependencies +RUN pip install -r requirements.txt + +# Install transformers +RUN pip install transformers + +# Modify app.py to enable share=True for Gradio +RUN sed -i 's/block\.launch()/block.launch(share=True)/g' app.py + +# Create pretrained models directory and download models +RUN mkdir -p pretrained && \ + cd pretrained && \ + wget https://huggingface.co/nvidia/PartPacker/resolve/main/vae.pt && \ + wget https://huggingface.co/nvidia/PartPacker/resolve/main/flow.pt + +# Set global library paths to ensure proper linking at runtime +ENV LD_LIBRARY_PATH="/workspace/miniconda3/envs/partpacker/lib:${LD_LIBRARY_PATH}" + +# Activate conda environment by default +RUN echo "conda activate partpacker" >> ~/.bashrc +SHELL ["/bin/bash", "--login", "-c"] + +# Cleanup +RUN apt-get clean && \ + rm -rf /var/lib/apt/lists/* && \ + conda clean -a -y + +# Expose port for Gradio app +EXPOSE 7860 + +# Set default command to bash +CMD ["/bin/bash"] \ No newline at end of file diff --git a/PartPacker/docker/README.md b/PartPacker/docker/README.md new file mode 100644 index 0000000..7aba14c --- /dev/null +++ b/PartPacker/docker/README.md @@ -0,0 +1,29 @@ +# Docker setup + +This docker setup is tested on Windows 10. + +make sure you are under this directory yourworkspace/PartPacker/docker + +Build docker image: + +``` +docker build -t partpacker:latest . +``` + +Run docker image at the first time: + +``` +docker run --name partpacker --gpus all -it -p 7860:7860 partpacker python app.py +``` + +After first time: +``` +docker start -a partpacker +``` + +Stop the container: +``` +docker stop partpacker +``` + +You can find the demo link showing in terminal, such as `https://94fc1ba77a08526e17.gradio.live/` or something similar else (it will be changed after each time to restart the container) to use the demo. \ No newline at end of file diff --git a/PartPacker/flow/__init__.py b/PartPacker/flow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/flow/configs/__init__.py b/PartPacker/flow/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/flow/configs/big_parts_strict_pvae.py b/PartPacker/flow/configs/big_parts_strict_pvae.py new file mode 100644 index 0000000..cbdfac5 --- /dev/null +++ b/PartPacker/flow/configs/big_parts_strict_pvae.py @@ -0,0 +1,33 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from .schema import ModelConfig + + +def make_config(): + + model_config = ModelConfig( + vae_conf="vae.configs.part_woenc", + vae_ckpt_path="pretrained/vae.pt", + qknorm=True, + qknorm_type="RMSNorm", + use_pos_embed=False, + dino_model="dinov2_vitg14", + hidden_dim=1536, + flow_shift=3.0, + logitnorm_mean=1.0, + logitnorm_std=1.0, + latent_size=4096, + use_parts=True, + ) + + return model_config diff --git a/PartPacker/flow/configs/schema.py b/PartPacker/flow/configs/schema.py new file mode 100644 index 0000000..7ec1db5 --- /dev/null +++ b/PartPacker/flow/configs/schema.py @@ -0,0 +1,57 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from typing import Literal, Optional + +import attrs + + +@attrs.define(slots=False) +class ModelConfig: + # vae + vae_conf: str = "vae.configs.part_woenc" + vae_ckpt_path: Optional[str] = None + + # learn & generate parts + use_parts: bool = False + part_embed_mode: Literal["element", "part", "part2_only"] = "part2_only" + shuffle_parts: bool = False + use_num_parts_cond: bool = False + + # flow matching hyper-params + flow_shift: float = 1.0 + logitnorm_mean: float = 0.0 + logitnorm_std: float = 1.0 + + # image encoder + dino_model: Literal["dinov2_vitl14_reg", "dinov2_vitg14"] = "dinov2_vitg14" + + # backbone DiT + hidden_dim: int = 1536 + num_heads: int = 16 + num_layers: int = 24 + qknorm: bool = True + qknorm_type: Literal["LayerNorm", "RMSNorm"] = "RMSNorm" + use_pos_embed: bool = False + + # latent code + latent_size: Optional[int] = None # if None, will load from vae + latent_dim: Optional[int] = None + + # preload vae weights + preload_vae: bool = True + + # preload dinov2 weights + preload_dinov2: bool = True + + # init weights from a pretrained checkpoint + pretrain_path: Optional[str] = None diff --git a/PartPacker/flow/flow_matching.py b/PartPacker/flow/flow_matching.py new file mode 100644 index 0000000..dc9f903 --- /dev/null +++ b/PartPacker/flow/flow_matching.py @@ -0,0 +1,58 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import numpy as np +import torch + + +class FlowMatchingScheduler: + def __init__(self, num_train_timesteps: int = 1000, shift: float = 1): + # set timesteps + self.num_train_timesteps = num_train_timesteps + self.shift = shift + + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.sigmas = sigmas # 1 --> 0 + self.timesteps = sigmas * num_train_timesteps # num_train_timesteps --> 1 + + # set device + def to(self, device): + self.sigmas = self.sigmas.to(device=device) + self.timesteps = self.timesteps.to(device=device) + + # add random noise to latent during training + def add_noise(self, latent: torch.Tensor, logit_mean: float = 1.0, logit_std: float = 1.0): + # latent: [B, ...] + # timesteps: [B] + # return: [B, ...] noisy_latent, [B, ...] noise, [B] timesteps + + # logit-normal sampling + u = torch.normal(mean=logit_mean, std=logit_std, size=(latent.shape[0],), device=self.sigmas.device) + u = torch.nn.functional.sigmoid(u) + + step_indices = (u * self.num_train_timesteps).long() + timesteps = self.timesteps[step_indices] + + sigmas = self.sigmas[step_indices].flatten() + + while len(sigmas.shape) < latent.ndim: + sigmas = sigmas.unsqueeze(-1) + + noise = torch.randn_like(latent) + noisy_latent = (1.0 - sigmas) * latent + sigmas * noise + + return noisy_latent, noise, timesteps diff --git a/PartPacker/flow/model.py b/PartPacker/flow/model.py new file mode 100644 index 0000000..e0b5032 --- /dev/null +++ b/PartPacker/flow/model.py @@ -0,0 +1,343 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import importlib + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import tqdm +from torchvision import transforms +from transformers import Dinov2Model + +from .configs.schema import ModelConfig +from .flow_matching import FlowMatchingScheduler +from .modules.dit import DiT +from ..vae.model import Model as VAE +from ..vae.utils import sync_timer + + +class Model(nn.Module): + def __init__(self, config: ModelConfig,device,dino_path,cpu_offload=False) -> None: + super().__init__() + self.config = config + self.precision = torch.bfloat16 + self.cpu_offload = cpu_offload + # image condition model (dinov2) + # if self.config.dino_model == "dinov2_vitg14": + # #self.dino = Dinov2Model.from_pretrained("facebook/dinov2-giant") + # elif self.config.dino_model == "dinov2_vitl14_reg": + # #self.dino = Dinov2Model.from_pretrained("facebook/dinov2-with-registers-large") + self.device = device + if dino_path: + self.dino = Dinov2Model.from_pretrained(dino_path) + else: + raise ValueError(f"DINOv2 model {self.config.dino_model} not supported") + + # hack to match our implementation + self.dino.layernorm = torch.nn.Identity() + + self.dino.eval().to(dtype=self.precision) + self.dino.requires_grad_(False) + + cond_dim = 1024 if self.config.dino_model == "dinov2_vitl14_reg" else 1536 + assert cond_dim == config.hidden_dim, "DINOv2 dim must match backbone dim" + + self.preprocess_cond_image = transforms.Compose( + [ + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) + + # vae encoder + vae_config = importlib.import_module(config.vae_conf).make_config() + self.vae = VAE(vae_config).eval().to(dtype=self.precision) + self.vae.requires_grad_(False) + + # load vae + if self.config.preload_vae: + try: + vae_ckpt = torch.load(self.config.vae_ckpt_path, weights_only=True) # local path + if "model" in vae_ckpt: + vae_ckpt = vae_ckpt["model"] + self.vae.load_state_dict(vae_ckpt, strict=True) + del vae_ckpt + print(f"Loaded VAE from {self.config.vae_ckpt_path}") + except Exception as e: + print( + f"Failed to load VAE from {self.config.vae_ckpt_path}: {e}, make sure you resumed from a valid checkpoint!" + ) + + # load info from vae config + if config.latent_size is None: + config.latent_size = self.vae.config.latent_size + if config.latent_dim is None: + config.latent_dim = self.vae.config.latent_dim + + # dit + self.dit = DiT( + hidden_dim=config.hidden_dim, + num_heads=config.num_heads, + num_layers=config.num_layers, + latent_size=config.latent_size, + latent_dim=config.latent_dim, + qknorm=config.qknorm, + qknorm_type=config.qknorm_type, + use_pos_embed=config.use_pos_embed, + use_parts=config.use_parts, + part_embed_mode=config.part_embed_mode, + ) + + # num_part condition + if self.config.use_num_parts_cond: + assert self.config.use_parts, "use_num_parts_cond requires use_parts" + self.num_part_embed = nn.Embedding(5, config.hidden_dim) + + # preload from a checkpoint (NOTE: this happens BEFORE checkpointer loading latest checkpoint!) + if self.config.pretrain_path is not None: + try: + ckpt = torch.load(self.config.pretrain_path) # local path + self.load_state_dict(ckpt["model"], strict=True) + del ckpt + print(f"Loaded DiT from {self.config.pretrain_path}") + except Exception as e: + print( + f"Failed to load DiT from {self.config.pretrain_path}: {e}, make sure you resumed from a valid checkpoint!" + ) + + # sampler + self.scheduler = FlowMatchingScheduler(shift=config.flow_shift) + + n_params = 0 + for p in self.dit.parameters(): + n_params += p.numel() + print(f"Number of parameters in DiT: {n_params/1e6:.2f}M") + + # override state_dict to exclude vae and dino, so we only save the trainable params. + def state_dict(self, *args, **kwargs): + state_dict = super().state_dict(*args, **kwargs) + + keys_to_del = [] + for k in state_dict.keys(): + if "vae" in k or "dino" in k: + keys_to_del.append(k) + + for k in keys_to_del: + del state_dict[k] + + return state_dict + + # override to support tolerant loading (only load matched shape) + def load_state_dict(self, state_dict, strict=True, assign=False): + local_state_dict = self.state_dict() + seen_keys = {k: False for k in local_state_dict.keys()} + for k, v in state_dict.items(): + if k in local_state_dict: + seen_keys[k] = True + if local_state_dict[k].shape == v.shape: + local_state_dict[k].copy_(v) + else: + print(f"mismatching shape for key {k}: loaded {local_state_dict[k].shape} but model has {v.shape}") + else: + print(f"unexpected key {k} in loaded state dict") + for k in seen_keys: + if not seen_keys[k]: + print(f"missing key {k} in loaded state dict") + + # this happens before checkpointer loading old models !!! + def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None: + super().on_train_start(memory_format=memory_format) + device = next(self.dit.parameters()).device + + self.dit.to(dtype=self.precision) + + if self.config.use_num_parts_cond: + self.num_part_embed.to(dtype=self.precision) + + # cast scheduler to device + self.scheduler.to(device) + + def get_cond(self, cond_image, num_part=None): + # image condition + cond_image = cond_image.to(dtype=self.precision) + with torch.no_grad(): + cond = self.dino(cond_image).last_hidden_state + cond = F.layer_norm(cond.float(), cond.shape[-1:]).to(dtype=self.precision) # [B, L, C] + + # num_part condition + if self.config.use_num_parts_cond: + if num_part is None: + # use a default value (2-10 parts) + num_part_coarse = torch.ones(cond.shape[0], dtype=torch.int64, device=cond.device) * 2 + else: + # coarse range + num_part_coarse = torch.ones(cond.shape[0], dtype=torch.int64, device=cond.device) + num_part_coarse[num_part == 2] = 1 + num_part_coarse[(num_part > 2) & (num_part <= 10)] = 2 + num_part_coarse[(num_part > 10) & (num_part <= 100)] = 3 + num_part_coarse[num_part > 100] = 4 + num_part_embed = self.num_part_embed(num_part_coarse).unsqueeze(1) # [B, 1, C] + cond = torch.cat([cond, num_part_embed], dim=1) # [B, L+1, C] + + return cond + + def training_step( + self, + data: dict[str, torch.Tensor], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + output = {} + loss = 0 + + cond_images = self.preprocess_cond_image( + data["cond_images"] + ) # [B, N, 3, 518, 518], we may load multiple (N) cond images for the same shape + B, N, C, H, W = cond_images.shape + + if self.config.use_num_parts_cond: + cond_num_part = data["num_part"].repeat_interleave(N, dim=0) + else: + cond_num_part = None + + cond = self.get_cond(cond_images.view(-1, C, H, W), cond_num_part) # [B*N, L, C] + + # random CFG dropout + if self.training: + mask = torch.rand((B * N, 1, 1), device=cond.device, dtype=cond.dtype) >= 0.1 + cond = cond * mask + + with torch.no_grad(): + # encode latent + if self.config.use_parts: + # encode two parts and concat latent + part0_data = {k.replace("_part0", ""): v for k, v in data.items() if "_part0" in k} + part1_data = {k.replace("_part1", ""): v for k, v in data.items() if "_part1" in k} + posterior0 = self.vae.encode(part0_data) + posterior1 = self.vae.encode(part1_data) + if self.training and self.config.shuffle_parts: + if np.random.rand() < 0.5: + posterior0, posterior1 = posterior1, posterior0 + latent = torch.cat( + [ + posterior0.mode().float().nan_to_num_(0), + posterior1.mode().float().nan_to_num_(0), + ], + dim=1, + ) # [B, 2L, C] + else: + posterior = self.vae.encode(data) + latent = posterior.mode().float().nan_to_num_(0) # use mean as the latent, [B, L, C] + + # repeat latent for each cond image + if N != 1: + latent = latent.repeat_interleave(N, dim=0) + + # random sample timesteps and add noise + noisy_latent, noise, timesteps = self.scheduler.add_noise( + latent, self.config.logitnorm_mean, self.config.logitnorm_std + ) + + noisy_latent = noisy_latent.to(dtype=self.precision) + model_pred = self.dit(noisy_latent, cond, timesteps) + + # flow-matching loss + target = noise - latent + loss = F.mse_loss(model_pred.float(), target.float()) + + # metrics + with torch.no_grad(): + output["scalar"] = {} # for wandb logging + output["scalar"]["loss_mse"] = loss.detach() + + return output, loss + + @torch.no_grad() + def validation_step( + self, + data: dict[str, torch.Tensor], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + return self.training_step(data, iteration) + + @torch.inference_mode() + @sync_timer("flow forward") + def forward( + self, + data: dict[str, torch.Tensor], + num_steps: int = 30, + cfg_scale: float = 7.0, + verbose: bool = True, + generator: torch.Generator | None = None, + ) -> dict[str, torch.Tensor]: + # the inference sampling + cond_images = self.preprocess_cond_image(data["cond_images"]) # [B, 3, 518, 518] + B = cond_images.shape[0] + assert B == 1, "Only support batch size 1 for now." + + # num_part condition + if self.config.use_num_parts_cond and "num_part" in data: + cond_num_part = data["num_part"] # [B,], int + else: + cond_num_part = None + if self.cpu_offload: + self.dino.to(device=self.device) + cond = self.get_cond(cond_images, cond_num_part) + if self.cpu_offload: + self.dino.to(device="cpu") + torch.cuda.empty_cache() + + if self.config.use_parts: + x = torch.randn( + B, + self.config.latent_size * 2, + self.config.latent_dim, + device=cond.device, + dtype=torch.float32, + generator=generator, + ) + else: + x = torch.randn( + B, + self.config.latent_size, + self.config.latent_dim, + device=cond.device, + dtype=torch.float32, + generator=generator, + ) + + cond_input = torch.cat([cond, torch.zeros_like(cond)], dim=0) + + # flow-matching + sigmas = np.linspace(1, 0, num_steps + 1) + sigmas = self.scheduler.shift * sigmas / (1 + (self.scheduler.shift - 1) * sigmas) + sigmas_pair = list((sigmas[i], sigmas[i + 1]) for i in range(num_steps)) + + for sigma, sigma_prev in tqdm.tqdm(sigmas_pair, desc="Flow Sampling", disable=not verbose): + # classifier-free guidance + timesteps = torch.tensor([1000 * sigma] * B * 2, device=x.device, dtype=x.dtype) + x_input = torch.cat([x, x], dim=0) + + # predict v + x_input = x_input.to(dtype=self.precision) + pred = self.dit(x_input, cond_input, timesteps).float() + cond_v, uncond_v = pred.chunk(2, dim=0) + pred_v = uncond_v + (cond_v - uncond_v) * cfg_scale + + # scheduler step + x = x - (sigma - sigma_prev) * pred_v + + output = {} + output["latent"] = x + + # leave mesh extraction to vae + return output diff --git a/PartPacker/flow/modules/__init__.py b/PartPacker/flow/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/flow/modules/dit.py b/PartPacker/flow/modules/dit.py new file mode 100644 index 0000000..fa70ce4 --- /dev/null +++ b/PartPacker/flow/modules/dit.py @@ -0,0 +1,235 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + +from ...vae.modules.attention import CrossAttention, SelfAttention + + +class FeedForward(nn.Module): + def __init__(self, dim, mult=4): + super().__init__() + self.net = nn.Sequential(nn.Linear(dim, dim * mult), nn.GELU(), nn.Linear(dim * mult, dim)) + + def forward(self, x): + return self.net(x) + + +# Adapted from https://github.com/facebookresearch/DiT/blob/main/models.py#L27 +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + dim: the dimension of the output. + max_period: controls the minimum frequency of the embeddings. + + Returns: + an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + dtype = next(self.mlp.parameters()).dtype # need to determine on the fly... + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_freq = t_freq.to(dtype=dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +class DiTLayer(nn.Module): + def __init__(self, dim, num_heads, qknorm=False, gradient_checkpointing=True, qknorm_type="LayerNorm"): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.gradient_checkpointing = gradient_checkpointing + + self.norm1 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.attn1 = SelfAttention(dim, num_heads, qknorm=qknorm, qknorm_type=qknorm_type) + self.norm2 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.attn2 = CrossAttention(dim, num_heads, context_dim=dim, qknorm=qknorm, qknorm_type=qknorm_type) + self.norm3 = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.ff = FeedForward(dim) + self.adaln_linear = nn.Linear(dim, dim * 6, bias=True) + + def forward(self, x, c, t_emb): + if self.training and self.gradient_checkpointing: + return checkpoint(self._forward, x, c, t_emb, use_reentrant=False) + else: + return self._forward(x, c, t_emb) + + def _forward(self, x, c, t_emb): + # x: [B, N, C], hidden states + # c: [B, M, C], condition (assume normed and projected to C) + # t_emb: [B, C], timestep embedding of adaln + # return: [B, N, C], updated hidden states + + B, N, C = x.shape + t_adaln = self.adaln_linear(F.silu(t_emb)).view(B, 6, -1) # [B, 6, C] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = t_adaln.chunk(6, dim=1) + + h = self.norm1(x) + h = h * (1 + scale_msa) + shift_msa + x = x + gate_msa * self.attn1(h) + + h = self.norm2(x) + x = x + self.attn2(h, c) + + h = self.norm3(x) + h = h * (1 + scale_mlp) + shift_mlp + x = x + gate_mlp * self.ff(h) + + return x + + +class DiT(nn.Module): + def __init__( + self, + hidden_dim=1024, + num_heads=16, + latent_size=2048, + latent_dim=8, + num_layers=24, + qknorm=False, + gradient_checkpointing=True, + qknorm_type="LayerNorm", + use_pos_embed=False, + use_parts=False, + part_embed_mode="part2_only", + ): + super().__init__() + + # project in + self.proj_in = nn.Linear(latent_dim, hidden_dim) + + # positional encoding (just use a learnable positional encoding) + self.use_pos_embed = use_pos_embed + if self.use_pos_embed: + self.pos_embed = nn.Parameter(torch.randn(1, latent_size, hidden_dim) / hidden_dim**0.5) + + # part encoding (a must to distinguish parts!) + self.use_parts = use_parts + self.part_embed_mode = part_embed_mode + if self.use_parts: + if self.part_embed_mode == "element": + self.part_embed = nn.Parameter(torch.randn(latent_size, hidden_dim) / hidden_dim**0.5) + elif self.part_embed_mode == "part": + self.part_embed = nn.Parameter(torch.randn(2, hidden_dim)) + elif self.part_embed_mode == "part2_only": + # we only add this to the second part to distinguish from the first part + self.part_embed = nn.Parameter(torch.randn(1, hidden_dim) / hidden_dim**0.5) + + # timestep encoding + self.timestep_embed = TimestepEmbedder(hidden_dim) + + # transformer layers + self.layers = nn.ModuleList( + [DiTLayer(hidden_dim, num_heads, qknorm, gradient_checkpointing, qknorm_type) for _ in range(num_layers)] + ) + + # project out + self.norm_out = nn.LayerNorm(hidden_dim, eps=1e-6, elementwise_affine=False) + self.proj_out = nn.Linear(hidden_dim, latent_dim) + + # init + self.init_weight() + + def init_weight(self): + # Initialize transformer layers + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.timestep_embed.mlp[0].weight, std=0.02) + nn.init.normal_(self.timestep_embed.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + for layer in self.layers: + nn.init.constant_(layer.adaln_linear.weight, 0) + nn.init.constant_(layer.adaln_linear.bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.proj_out.weight, 0) + nn.init.constant_(self.proj_out.bias, 0) + + def forward(self, x, c, t): + # x: [B, N, C], hidden states + # c: [B, M, C], condition (assume normed and projected to C) + # t: [B,], timestep + # return: [B, N, C], updated hidden states + + B, N, C = x.shape + + # project in + x = self.proj_in(x) + + # positional encoding + if self.use_pos_embed: + x = x + self.pos_embed + + # part encoding + if self.use_parts: + if self.part_embed_mode == "element": + x += self.part_embed + elif self.part_embed_mode == "part": + x[:, : x.shape[1] // 2, :] += self.part_embed[0] + x[:, x.shape[1] // 2 :, :] += self.part_embed[1] + elif self.part_embed_mode == "part2_only": + x[:, x.shape[1] // 2 :, :] += self.part_embed[0] + + # timestep encoding + t_emb = self.timestep_embed(t) # [B, C] + + # transformer layers + for layer in self.layers: + x = layer(x, c, t_emb) + + # project out + x = self.norm_out(x) + x = self.proj_out(x) + + return x diff --git a/PartPacker/flow/scripts/infer.py b/PartPacker/flow/scripts/infer.py new file mode 100644 index 0000000..351b5cb --- /dev/null +++ b/PartPacker/flow/scripts/infer.py @@ -0,0 +1,184 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import argparse +import glob +import importlib +import os +from datetime import datetime + +import cv2 +import kiui +import numpy as np +import rembg +import torch +import trimesh + +from ..model import Model +from ..utils import get_random_color, recenter_foreground +from ...vae.utils import postprocess_mesh + +# PYTHONPATH=. python flow/scripts/infer.py +parser = argparse.ArgumentParser() +parser.add_argument( + "--config", + type=str, + help="config file path", + default="flow.configs.big_parts_strict_pvae", +) +parser.add_argument( + "--ckpt_path", + type=str, + help="checkpoint path", + default="pretrained/flow.pt", +) +parser.add_argument("--input", type=str, help="input directory", default="assets/images/") +parser.add_argument("--limit", type=int, help="limit number of images", default=-1) +parser.add_argument("--output_dir", type=str, help="output directory", default="output/") +parser.add_argument("--grid_res", type=int, help="grid resolution", default=384) +parser.add_argument("--num_steps", type=int, help="number of cfg steps", default=50) +parser.add_argument("--cfg_scale", type=float, help="cfg scale", default=7.0) +parser.add_argument("--num_repeats", type=int, help="number of repeats per image", default=1) +parser.add_argument("--num_faces", type=int, help="target number of faces for decimation", default=-1) +parser.add_argument("--seed", type=int, help="seed", default=42) +args = parser.parse_args() + +TRIMESH_GLB_EXPORT = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).astype(np.float32) + +bg_remover = rembg.new_session() + + +def preprocess_image(path): + input_image = kiui.read_image(path, mode="uint8", order="RGBA") + + # bg removal if there is no alpha channel + if input_image.shape[-1] == 3: + input_image = rembg.remove(input_image, session=bg_remover) # [H, W, 4] + + mask = input_image[..., -1] > 0 + image = recenter_foreground(input_image, mask, border_ratio=0.1) + image = cv2.resize(image, (518, 518), interpolation=cv2.INTER_LINEAR) + image = image.astype(np.float32) / 255.0 + image = image[..., :3] * image[..., 3:4] + (1 - image[..., 3:4]) # white background + return image + + +print(f"Loading checkpoint from {args.ckpt_path}") +ckpt_dict = torch.load(args.ckpt_path, weights_only=True) + +# delete all keys other than model +if "model" in ckpt_dict: + ckpt_dict = ckpt_dict["model"] + +# instantiate model +print(f"Instantiating model from {args.config}") +model_config = importlib.import_module(args.config).make_config() +model = Model(model_config).eval().cuda().bfloat16() + +# load weight +print(f"Loading weights from {args.ckpt_path}") +model.load_state_dict(ckpt_dict, strict=True) + +# output folder +timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") +workspace = os.path.join(args.output_dir, "flow_" + args.config.split(".")[-1] + "_" + timestamp) +if not os.path.exists(workspace): + os.makedirs(workspace) +else: + os.system(f"rm {workspace}/*") +print(f"Output directory: {workspace}") + +# load test images +if os.path.isdir(args.input): + paths = glob.glob(os.path.join(args.input, "*")) + paths = sorted(paths) + if args.limit > 0: + paths = paths[: args.limit] +else: # single file + paths = [args.input] + +for path in paths: + name = os.path.splitext(os.path.basename(path))[0] + print(f"Processing {name}") + + image = preprocess_image(path) + + kiui.write_image(os.path.join(workspace, name + ".jpg"), image) + image = torch.from_numpy(image).permute(2, 0, 1).contiguous().unsqueeze(0).float().cuda() + + # run model + data = {"cond_images": image} + + for i in range(args.num_repeats): + + kiui.seed_everything(args.seed + i) + + with torch.inference_mode(): + results = model(data, num_steps=args.num_steps, cfg_scale=args.cfg_scale) + + latent = results["latent"] + # kiui.lo(latent) + + # query mesh + if model.config.use_parts: + data_part0 = {"latent": latent[:, : model.config.latent_size, :]} + data_part1 = {"latent": latent[:, model.config.latent_size :, :]} + + with torch.inference_mode(): + results_part0 = model.vae(data_part0, resolution=args.grid_res) + results_part1 = model.vae(data_part1, resolution=args.grid_res) + + vertices, faces = results_part0["meshes"][0] + mesh_part0 = trimesh.Trimesh(vertices, faces) + mesh_part0.vertices = mesh_part0.vertices @ TRIMESH_GLB_EXPORT.T + mesh_part0 = postprocess_mesh(mesh_part0, args.num_faces) + parts = mesh_part0.split(only_watertight=False) + + vertices, faces = results_part1["meshes"][0] + mesh_part1 = trimesh.Trimesh(vertices, faces) + mesh_part1.vertices = mesh_part1.vertices @ TRIMESH_GLB_EXPORT.T + mesh_part1 = postprocess_mesh(mesh_part1, args.num_faces) + parts.extend(mesh_part1.split(only_watertight=False)) + + # some parts only have 1 face, seems a problem of trimesh.split. + parts = [part for part in parts if len(part.faces) > 10] + + # split connected components and assign different colors + for j, part in enumerate(parts): + # each component uses a random color + part.visual.vertex_colors = get_random_color(j, use_float=True) + + mesh = trimesh.Scene(parts) + # export the whole mesh + mesh.export(os.path.join(workspace, name + "_" + str(i) + ".glb")) + + # export each part + for j, part in enumerate(parts): + part.export(os.path.join(workspace, name + "_" + str(i) + "_part" + str(j) + ".glb")) + + # export dual volumes + mesh_part0.export(os.path.join(workspace, name + "_" + str(i) + "_vol0.glb")) + mesh_part1.export(os.path.join(workspace, name + "_" + str(i) + "_vol1.glb")) + + else: + data = {"latent": latent} + + with torch.inference_mode(): + results = model.vae(data, resolution=args.grid_res) + + vertices, faces = results["meshes"][0] + mesh = trimesh.Trimesh(vertices, faces) + mesh = postprocess_mesh(mesh, args.num_faces) + + # kiui.lo(mesh.vertices, mesh.faces) + mesh.vertices = mesh.vertices @ TRIMESH_GLB_EXPORT.T + mesh.export(os.path.join(workspace, name + "_" + str(i) + ".glb")) diff --git a/PartPacker/flow/utils.py b/PartPacker/flow/utils.py new file mode 100644 index 0000000..019fcf0 --- /dev/null +++ b/PartPacker/flow/utils.py @@ -0,0 +1,119 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from typing import Optional + +import cv2 +import numpy as np + + +def recenter_foreground(image, mask, border_ratio: float = 0.1): + """recenter an image to leave some empty space at the image border. + + Args: + image (ndarray): input image, float/uint8 [H, W, 3/4] + mask (ndarray): alpha mask, bool [H, W] + border_ratio (float, optional): border ratio, image will be resized to (1 - border_ratio). Defaults to 0.1. + + Returns: + ndarray: output image, float/uint8 [H, W, 3/4] + """ + + # empty foreground: just return + if mask.sum() == 0: + return image + + return_int = False + if image.dtype == np.uint8: + image = image.astype(np.float32) / 255 + return_int = True + + H, W, C = image.shape + size = max(H, W) + + # default to white bg if rgb, but use 0 if rgba + if C == 3: + result = np.ones((size, size, C), dtype=np.float32) + else: + result = np.zeros((size, size, C), dtype=np.float32) + + coords = np.nonzero(mask) + x_min, x_max = coords[0].min(), coords[0].max() + y_min, y_max = coords[1].min(), coords[1].max() + h = x_max - x_min + w = y_max - y_min + desired_size = int(size * (1 - border_ratio)) + scale = desired_size / max(h, w) + h2 = int(h * scale) + w2 = int(w * scale) + x2_min = (size - h2) // 2 + x2_max = x2_min + h2 + y2_min = (size - w2) // 2 + y2_max = y2_min + w2 + result[x2_min:x2_max, y2_min:y2_max] = cv2.resize( + image[x_min:x_max, y_min:y_max], (w2, h2), interpolation=cv2.INTER_AREA + ) + + if return_int: + result = (result * 255).astype(np.uint8) + + return result + + +def get_random_color(index: Optional[int] = None, use_float: bool = False): + # some pleasing colors + # matplotlib.colormaps['Set3'].colors + matplotlib.colormaps['Set2'].colors + matplotlib.colormaps['Set1'].colors + palette = np.array( + [ + [141, 211, 199, 255], + [255, 255, 179, 255], + [190, 186, 218, 255], + [251, 128, 114, 255], + [128, 177, 211, 255], + [253, 180, 98, 255], + [179, 222, 105, 255], + [252, 205, 229, 255], + [217, 217, 217, 255], + [188, 128, 189, 255], + [204, 235, 197, 255], + [255, 237, 111, 255], + [102, 194, 165, 255], + [252, 141, 98, 255], + [141, 160, 203, 255], + [231, 138, 195, 255], + [166, 216, 84, 255], + [255, 217, 47, 255], + [229, 196, 148, 255], + [179, 179, 179, 255], + [228, 26, 28, 255], + [55, 126, 184, 255], + [77, 175, 74, 255], + [152, 78, 163, 255], + [255, 127, 0, 255], + [255, 255, 51, 255], + [166, 86, 40, 255], + [247, 129, 191, 255], + [153, 153, 153, 255], + ], + dtype=np.uint8, + ) + + if index is None: + index = np.random.randint(0, len(palette)) + + if index >= len(palette): + index = index % len(palette) + + if use_float: + return palette[index].astype(np.float32) / 255 + else: + return palette[index] diff --git a/PartPacker/requirements.lock.txt b/PartPacker/requirements.lock.txt new file mode 100644 index 0000000..36054a7 --- /dev/null +++ b/PartPacker/requirements.lock.txt @@ -0,0 +1,15 @@ +torch == 2.5.1+cu121 +numpy == 1.26.4 +trimesh == 4.6.9 +fpsample == 0.3.3 +einops == 0.8.0 +onnxruntime == 1.18.0 +rembg == 2.0.60 +kiui == 0.2.15 +torchvision == 0.20.1+cu121 +pymcubes == 0.1.6 +tqdm == 4.67.0 +opencv-python == 4.10.0.84 +pymeshlab == 2023.12.post2 +gradio == 5.25.2 +transformers == 4.46.2 diff --git a/PartPacker/vae/__init__.py b/PartPacker/vae/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/vae/configs/__init__.py b/PartPacker/vae/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/vae/configs/part_woenc.py b/PartPacker/vae/configs/part_woenc.py new file mode 100644 index 0000000..5c4cb46 --- /dev/null +++ b/PartPacker/vae/configs/part_woenc.py @@ -0,0 +1,30 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from .schema import ModelConfig + + +def make_config(): + + model_config = ModelConfig( + use_salient_point=True, + latent_size=4096, + cutoff_fps_point=(256, 512, 512, 512, 1024, 1024, 2048), + cutoff_fps_salient_point=(0, 0, 256, 512, 512, 1024, 2048), + cutoff_fps_prob=(0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 0.2), + kl_weight=1e-3, + salient_attn_mode="dual", + num_enc_layers=0, + num_dec_layers=24, + ) + + return model_config diff --git a/PartPacker/vae/configs/schema.py b/PartPacker/vae/configs/schema.py new file mode 100644 index 0000000..15e9f8c --- /dev/null +++ b/PartPacker/vae/configs/schema.py @@ -0,0 +1,55 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from typing import Literal, Optional, Tuple + +import attrs + + +@attrs.define(slots=False) +class ModelConfig: + # input + use_salient_point: bool = True + + # random cutoff during training + cutoff_fps_point: Tuple[int, ...] = (256, 512, 512, 512, 1024, 1024, 2048) + cutoff_fps_salient_point: Tuple[int, ...] = (0, 0, 256, 512, 512, 1024, 2048) + cutoff_fps_prob: Tuple[float, ...] = (0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 0.2) # sum to 1.0 + + # backbone transformer + num_enc_layers: int = 0 + hidden_dim: int = 1024 + num_heads: int = 16 + num_dec_layers: int = 24 + dec_hidden_dim: int = 1024 + dec_num_heads: int = 16 + qknorm: bool = True + qknorm_type: Literal["LayerNorm", "RMSNorm"] = "LayerNorm" # type of qknorm + salient_attn_mode: Literal["dual_shared", "single", "dual"] = "dual" + + # query decoder + fourier_version: Literal["v1", "v2", "v3"] = "v3" + point_fourier_dim: int = 48 # must be divisible by 6 (sin/cos, x/y/z) + query_hidden_dim: int = 1024 + query_num_heads: int = 16 + use_flash_query: bool = False + + # latent code + latent_size: int = 4096 # == num_fps_point + num_fps_salient_point + latent_dim: int = 64 + + # loss + use_ae: bool = False # if true, variance will be ignored, and kl_weight is used as a L2 norm weight + kl_weight: float = 1e-3 + + # init weights from a pretrained checkpoint + pretrain_path: Optional[str] = None diff --git a/PartPacker/vae/model.py b/PartPacker/vae/model.py new file mode 100644 index 0000000..cc4598d --- /dev/null +++ b/PartPacker/vae/model.py @@ -0,0 +1,451 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +from typing import Literal + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .configs.schema import ModelConfig +from .modules.transformer import AttentionBlock, FlashQueryLayer +from .utils import ( + DiagonalGaussianDistribution, + DummyLatent, + calculate_iou, + calculate_metrics, + construct_grid_points, + extract_mesh, + sync_timer, +) + + +class Model(nn.Module): + def __init__(self, config: ModelConfig) -> None: + super().__init__() + self.config = config + + self.precision = torch.bfloat16 # manually handle low-precision training, always use bf16 + + # point encoder + self.proj_input = nn.Linear(3 + config.point_fourier_dim, config.hidden_dim) + + self.perceiver = AttentionBlock( + config.hidden_dim, + num_heads=config.num_heads, + dim_context=config.hidden_dim, + qknorm=config.qknorm, + qknorm_type=config.qknorm_type, + ) + + if self.config.salient_attn_mode == "dual": + self.perceiver_dorases = AttentionBlock( + config.hidden_dim, + num_heads=config.num_heads, + dim_context=config.hidden_dim, + qknorm=config.qknorm, + qknorm_type=config.qknorm_type, + ) + + # self-attention encoder + self.encoder = nn.ModuleList( + [ + AttentionBlock( + config.hidden_dim, config.num_heads, qknorm=config.qknorm, qknorm_type=config.qknorm_type + ) + for _ in range(config.num_enc_layers) + ] + ) + + # vae bottleneck + self.norm_down = nn.LayerNorm(config.hidden_dim) + self.proj_down_mean = nn.Linear(config.hidden_dim, config.latent_dim) + if not self.config.use_ae: + self.proj_down_std = nn.Linear(config.hidden_dim, config.latent_dim) + self.proj_up = nn.Linear(config.latent_dim, config.dec_hidden_dim) + + # self-attention decoder + self.decoder = nn.ModuleList( + [ + AttentionBlock( + config.dec_hidden_dim, config.dec_num_heads, qknorm=config.qknorm, qknorm_type=config.qknorm_type + ) + for _ in range(config.num_dec_layers) + ] + ) + + # cross-attention query + self.proj_query = nn.Linear(3 + config.point_fourier_dim, config.query_hidden_dim) + if self.config.use_flash_query: + self.norm_query_context = nn.LayerNorm(config.hidden_dim, eps=1e-6, elementwise_affine=False) + self.attn_query = FlashQueryLayer( + config.query_hidden_dim, + num_heads=config.query_num_heads, + dim_context=config.hidden_dim, + qknorm=config.qknorm, + qknorm_type=config.qknorm_type, + ) + else: + self.attn_query = AttentionBlock( + config.query_hidden_dim, + num_heads=config.query_num_heads, + dim_context=config.hidden_dim, + qknorm=config.qknorm, + qknorm_type=config.qknorm_type, + ) + self.norm_out = nn.LayerNorm(config.query_hidden_dim) + self.proj_out = nn.Linear(config.query_hidden_dim, 1) + + # preload from a checkpoint (NOTE: this happens BEFORE checkpointer loading latest checkpoint!) + if self.config.pretrain_path is not None: + try: + ckpt = torch.load(self.config.pretrain_path) # local path + self.load_state_dict(ckpt["model"], strict=True) + del ckpt + print(f"Loaded VAE from {self.config.pretrain_path}") + except Exception as e: + print( + f"Failed to load VAE from {self.config.pretrain_path}: {e}, make sure you resumed from a valid checkpoint!" + ) + + # log + n_params = 0 + for p in self.parameters(): + n_params += p.numel() + print(f"Number of parameters in VAE: {n_params / 1e6:.2f}M") + + # override to support tolerant loading (only load matched shape) + def load_state_dict(self, state_dict, strict=True, assign=False): + local_state_dict = self.state_dict() + seen_keys = {k: False for k in local_state_dict.keys()} + for k, v in state_dict.items(): + if k in local_state_dict: + seen_keys[k] = True + if local_state_dict[k].shape == v.shape: + local_state_dict[k].copy_(v) + else: + print(f"mismatching shape for key {k}: loaded {local_state_dict[k].shape} but model has {v.shape}") + else: + print(f"unexpected key {k} in loaded state dict") + for k in seen_keys: + if not seen_keys[k]: + print(f"missing key {k} in loaded state dict") + + def fourier_encoding(self, points: torch.Tensor): + # points: [B, N, 3], float32 for precision + # assert points.dtype == torch.float32, "Query points must be float32" + + F = self.config.point_fourier_dim // (2 * points.shape[-1]) + + if self.config.fourier_version == "v1": # default + exponent = torch.arange(1, F + 1, device=points.device, dtype=torch.float32) / F # [F], range from 0 to 1 + freq_band = 512**exponent # [F], min frequency is 1, max frequency is 1/freq + freq_band *= torch.pi + elif self.config.fourier_version == "v2": + exponent = torch.arange(F, device=points.device, dtype=torch.float32) / (F - 1) # [F], range from 0 to 1 + freq_band = 1024**exponent # [F] + freq_band *= torch.pi + elif self.config.fourier_version == "v3": # hunyuan3d-2 + freq_band = 2 ** torch.arange(F, device=points.device, dtype=torch.float32) # [F] + + spectrum = points.unsqueeze(-1) * freq_band # [B,...,3,F] + sin, cos = spectrum.sin(), spectrum.cos() # [B,...,3,F] + input_enc = torch.stack([sin, cos], dim=-2) # [B,...,3,2,F] + input_enc = input_enc.view(*points.shape[:-1], -1) # [B,...,6F] = [B,...,dim] + return torch.cat([input_enc, points], dim=-1).to(dtype=self.precision) # [B,...,dim+input_dim] + + def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None: + super().on_train_start(memory_format=memory_format) + self.to(dtype=self.precision, memory_format=memory_format) # use bfloat16 for training + + def encode(self, data: dict[str, torch.Tensor]): + # uniform points + pointcloud = data["pointcloud"] # [B, N, 3] + + # fourier embed and project + pointcloud = self.fourier_encoding(pointcloud) # [B, N, 3+C] + pointcloud = self.proj_input(pointcloud) # [B, N, hidden_dim] + + # salient points + if self.config.use_salient_point: + pointcloud_dorases = data["pointcloud_dorases"] # [B, M, 3] + + # fourier embed and project (shared weights) + pointcloud_dorases = self.fourier_encoding(pointcloud_dorases) # [B, M, 3+C] + pointcloud_dorases = self.proj_input(pointcloud_dorases) # [B, M, hidden_dim] + + # gather fps point + fps_indices = data["fps_indices"] # [B, N'] + pointcloud_query = torch.gather(pointcloud, 1, fps_indices.unsqueeze(-1).expand(-1, -1, pointcloud.shape[-1])) + + if self.config.use_salient_point: + fps_indices_dorases = data["fps_indices_dorases"] # [B, M'] + + if fps_indices_dorases.shape[1] > 0: + pointcloud_query_dorases = torch.gather( + pointcloud_dorases, + 1, + fps_indices_dorases.unsqueeze(-1).expand(-1, -1, pointcloud_dorases.shape[-1]), + ) + + # combine both fps points as the query + pointcloud_query = torch.cat( + [pointcloud_query, pointcloud_query_dorases], dim=1 + ) # [B, N'+M', hidden_dim] + + # dual cross-attention + if self.config.salient_attn_mode == "dual_shared": + hidden_states = self.perceiver(pointcloud_query, pointcloud) + self.perceiver( + pointcloud_query, pointcloud_dorases + ) # [B, N'+M', hidden_dim] + elif self.config.salient_attn_mode == "dual": + hidden_states = self.perceiver(pointcloud_query, pointcloud) + self.perceiver_dorases( + pointcloud_query, pointcloud_dorases + ) + else: # single, hunyuan3d-2 style + hidden_states = self.perceiver(pointcloud_query, torch.cat([pointcloud, pointcloud_dorases], dim=1)) + else: + hidden_states = self.perceiver(pointcloud_query, pointcloud) # [B, N', hidden_dim] + + # encoder + for block in self.encoder: + hidden_states = block(hidden_states) + + # bottleneck + hidden_states = self.norm_down(hidden_states) + latent_mean = self.proj_down_mean(hidden_states).float() + if not self.config.use_ae: + latent_std = self.proj_down_std(hidden_states).float() + posterior = DiagonalGaussianDistribution(latent_mean, latent_std) + else: + posterior = DummyLatent(latent_mean) + + return posterior + + def decode(self, latent: torch.Tensor): + latent = latent.to(dtype=self.precision) + hidden_states = self.proj_up(latent) + + for block in self.decoder: + hidden_states = block(hidden_states) + + return hidden_states + + def query(self, query_points: torch.Tensor, hidden_states: torch.Tensor): + # query_points: [B, N, 3], float32 to keep the precision + + query_points = self.fourier_encoding(query_points) # [B, N, 3+C] + query_points = self.proj_query(query_points) # [B, N, hidden_dim] + + # cross attention + query_output = self.attn_query(query_points, hidden_states) # [B, N, hidden_dim] + + # output linear + query_output = self.norm_out(query_output) + pred = self.proj_out(query_output) # [B, N, 1] + + return pred + + def training_step( + self, + data: dict[str, torch.Tensor], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + output = {} + + # cut off fps point during training for progressive flow + if self.training: + # randomly choose from a set of cutoff candidates + cutoff_index = np.random.choice(len(self.config.cutoff_fps_prob), p=self.config.cutoff_fps_prob) + cutoff_fps_point = self.config.cutoff_fps_point[cutoff_index] + cutoff_fps_salient_point = self.config.cutoff_fps_salient_point[cutoff_index] + # prefix of FPS points are still FPS points + data["fps_indices"] = data["fps_indices"][:, :cutoff_fps_point] + if self.config.use_salient_point: + data["fps_indices_dorases"] = data["fps_indices_dorases"][:, :cutoff_fps_salient_point] + + loss = 0 + + # encode + posterior = self.encode(data) + latent_geom = posterior.sample() if self.training else posterior.mode() + + # decode + hidden_states = self.decode(latent_geom) + + # cross-attention query + query_points = data["query_points"] # [B, N, 3], float32 + + # the context norm can be moved out to avoid repeated computation + if self.config.use_flash_query: + hidden_states = self.norm_query_context(hidden_states) + + pred = self.query(query_points, hidden_states).squeeze(-1).float() # [B, N] + gt = data["query_gt"].float() # [B, N], in [-1, 1] + + # main loss + loss_mse = F.mse_loss(pred, gt, reduction="mean") + loss += loss_mse + + loss_l1 = F.l1_loss(pred, gt, reduction="mean") + loss += loss_l1 + + # kl loss + loss_kl = posterior.kl().mean() + loss += self.config.kl_weight * loss_kl + + # metrics + with torch.no_grad(): + output["scalar"] = {} # for wandb logging + output["scalar"]["loss_mse"] = loss_mse.detach() + output["scalar"]["loss_l1"] = loss_l1.detach() + output["scalar"]["loss_kl"] = loss_kl.detach() + output["scalar"]["iou_fg"] = calculate_iou(pred, gt, target_value=1) + output["scalar"]["iou_bg"] = calculate_iou(pred, gt, target_value=0) + output["scalar"]["precision"], output["scalar"]["recall"], output["scalar"]["f1"] = calculate_metrics( + pred, gt, target_value=1 + ) + + return output, loss + + @torch.no_grad() + def validation_step( + self, + data: dict[str, torch.Tensor], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + return self.training_step(data, iteration) + + @torch.inference_mode() + @sync_timer("vae forward") + def forward( + self, + data: dict[str, torch.Tensor], + mode: Literal["dense", "hierarchical"] = "hierarchical", + max_samples_per_iter: int = 512**2, + resolution: int = 512, + min_resolution: int = 64, # for hierarchical + ) -> dict[str, torch.Tensor]: + output = {} + + # encode + if "latent" in data: + latent = data["latent"] + else: + posterior = self.encode(data) + output["posterior"] = posterior + latent = posterior.mode() + + output["latent"] = latent + B = latent.shape[0] + + # decode + hidden_states = self.decode(latent) + output["hidden_states"] = hidden_states # [B, N, hidden_dim] for the last cross-attention decoder + + # the context norm can be moved out to avoid repeated computation + if self.config.use_flash_query: + hidden_states = self.norm_query_context(hidden_states) + + # query + def chunked_query(grid_points): + if grid_points.shape[0] <= max_samples_per_iter: + return self.query(grid_points.unsqueeze(0), hidden_states).squeeze(-1) # [B, N] + all_pred = [] + for i in range(0, grid_points.shape[0], max_samples_per_iter): + grid_chunk = grid_points[i : i + max_samples_per_iter] + pred_chunk = self.query(grid_chunk.unsqueeze(0), hidden_states) + all_pred.append(pred_chunk) + return torch.cat(all_pred, dim=1).squeeze(-1) # [B, N] + + if mode == "dense": + grid_points = construct_grid_points(resolution).to(latent.device) + grid_points = grid_points.contiguous().view(-1, 3) + grid_vals = chunked_query(grid_points).float().view(B, resolution + 1, resolution + 1, resolution + 1) + + elif mode == "hierarchical": + assert resolution >= min_resolution, "Resolution must be greater than or equal to min_resolution" + assert B == 1, "Only one batch is supported for hierarchical mode" + + resolutions = [] + res = resolution + while res >= min_resolution: + resolutions.append(res) + res = res // 2 + resolutions.reverse() # e.g., [64, 128, 256, 512] + + # dense-query the coarsest resolution + res = resolutions[0] + grid_points = construct_grid_points(res).to(latent.device) + grid_points = grid_points.contiguous().view(-1, 3) + grid_vals = chunked_query(grid_points).float().view(res + 1, res + 1, res + 1) + + # sparse-query finer resolutions + dilate_kernel_3 = torch.ones(1, 1, 3, 3, 3, dtype=torch.float32, device=latent.device) + dilate_kernel_5 = torch.ones(1, 1, 5, 5, 5, dtype=torch.float32, device=latent.device) + for i in range(1, len(resolutions)): + res = resolutions[i] + # get the boundary grid mask in the coarser grid (where the grid_vals have different signs with at least one of its neighbors) + grid_signs = grid_vals >= 0 + mask = torch.zeros_like(grid_signs) + mask[1:, :, :] += grid_signs[1:, :, :] != grid_signs[:-1, :, :] + mask[:-1, :, :] += grid_signs[:-1, :, :] != grid_signs[1:, :, :] + mask[:, 1:, :] += grid_signs[:, 1:, :] != grid_signs[:, :-1, :] + mask[:, :-1, :] += grid_signs[:, :-1, :] != grid_signs[:, 1:, :] + mask[:, :, 1:] += grid_signs[:, :, 1:] != grid_signs[:, :, :-1] + mask[:, :, :-1] += grid_signs[:, :, :-1] != grid_signs[:, :, 1:] + # empirical: also add those with abs(grid_vals) < 0.95 + mask += grid_vals.abs() < 0.95 + mask = (mask > 0).float() + # empirical: dilate the coarse mask + if res < 512: + mask = mask.unsqueeze(0).unsqueeze(0) + mask = F.conv3d(mask, weight=dilate_kernel_3, padding=1) + mask = mask.squeeze(0).squeeze(0) + # get the coarse coordinates + cidx_x, cidx_y, cidx_z = torch.nonzero(mask, as_tuple=True) + # fill to the fine indices + mask_fine = torch.zeros(res + 1, res + 1, res + 1, dtype=torch.float32, device=latent.device) + mask_fine[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1 + # empirical: dilate the fine mask + if res < 512: + mask_fine = mask_fine.unsqueeze(0).unsqueeze(0) + mask_fine = F.conv3d(mask_fine, weight=dilate_kernel_3, padding=1) + mask_fine = mask_fine.squeeze(0).squeeze(0) + else: + mask_fine = mask_fine.unsqueeze(0).unsqueeze(0) + mask_fine = F.conv3d(mask_fine, weight=dilate_kernel_5, padding=2) + mask_fine = mask_fine.squeeze(0).squeeze(0) + # get the fine coordinates + fidx_x, fidx_y, fidx_z = torch.nonzero(mask_fine, as_tuple=True) + # convert to float query points + query_points = torch.stack([fidx_x, fidx_y, fidx_z], dim=-1) # [N, 3] + query_points = query_points * 2 / res - 1 # [N, 3], in [-1, 1] + # query + pred = chunked_query(query_points).float() + # fill to the fine indices + grid_vals = torch.full((res + 1, res + 1, res + 1), -100.0, dtype=torch.float32, device=latent.device) + grid_vals[fidx_x, fidx_y, fidx_z] = pred + # print(f"[INFO] hierarchical: resolution: {res}, valid coarse points: {len(cidx_x)}, valid fine points: {len(fidx_x)}") + + grid_vals = grid_vals.unsqueeze(0) # [1, res+1, res+1, res+1] + grid_vals[grid_vals <= -100.0] = float("nan") # use nans to ignore invalid regions + + # extract mesh + meshes = [] + for b in range(B): + vertices, faces = extract_mesh(grid_vals[b], resolution) + meshes.append((vertices, faces)) + output["meshes"] = meshes + + return output diff --git a/PartPacker/vae/modules/__init__.py b/PartPacker/vae/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PartPacker/vae/modules/attention.py b/PartPacker/vae/modules/attention.py new file mode 100644 index 0000000..670f823 --- /dev/null +++ b/PartPacker/vae/modules/attention.py @@ -0,0 +1,261 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import ( # , unpad_input # noqa + index_first_axis, + pad_input, + ) + + FLASH_ATTN_AVAILABLE = True +except Exception as e: + print("[WARN] flash_attn not available, using torch/naive implementation") + FLASH_ATTN_AVAILABLE = False + + +# Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py#L98 +# flashattn 2.7.0 changes the API, we are overriding it here +def unpad_input(hidden_states, attention_mask): + """ + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + """ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the + # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim + # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to + # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, + # so we write custom forward and backward to make it a bit faster. + return ( + index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def attention(q, k, v, mask_q=None, mask_kv=None, dropout=0, causal=False, window_size=(-1, -1), backend="torch"): + # q: (B, N, H, D) + # k: (B, M, H, D) + # v: (B, M, H, D) + # mask_q: (B, N) + # mask_kv: (B, M) + # return: (B, N, H, D) + + B, N, H, D = q.shape + M = k.shape[1] + + if causal: + assert N == 1 or N == M, "Causal mask only supports self-attention" + + # unmasked case (usually inference) + # will ignore window_size except flash-attn impl. Only provide the effective window! + if mask_q is None and mask_kv is None: + if backend == "flash-attn" and FLASH_ATTN_AVAILABLE: + return flash_attn_func(q, k, v, dropout, causal=causal, window_size=window_size) # [B, N, H, D] + elif backend == "torch": # torch implementation + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=dropout, is_causal=causal) + out = out.permute(0, 2, 1, 3).contiguous() + return out + else: # naive implementation + q = q.transpose(1, 2).reshape(B * H, N, D) + k = k.transpose(1, 2).reshape(B * H, M, D) + v = v.transpose(1, 2).reshape(B * H, M, D) + w = torch.bmm(q, k.transpose(1, 2)) / (D**0.5) # [B*H, N, M] + if causal and N > 1: + causal_mask = torch.full((N, M), float("-inf"), device=w.device, dtype=w.dtype) + causal_mask = torch.triu(causal_mask, diagonal=1) + w = w + causal_mask.unsqueeze(0) + w = F.softmax(w, dim=-1) + if dropout > 0: + w = F.dropout(w, p=dropout) + out = torch.bmm(w, v) # [B*H, N, D] + out = out.reshape(B, H, N, D).transpose(1, 2).contiguous() # [B, N, H, D] + return out + + # at least one of q or kv is masked (training) + # only support flash-attn for now... + if mask_q is None: + mask_q = torch.ones(B, N, dtype=torch.bool, device=q.device) + elif mask_kv is None: + mask_kv = torch.ones(B, M, dtype=torch.bool, device=q.device) + + if FLASH_ATTN_AVAILABLE: + # unpad (gather) input + # mask_q: [B, N], first row has N1 1s, second row has N2 1s, ... + # indices: [Ns,], Ns = N1 + N2 + ... + # cu_seqlens_q: [B+1,], (0, N1, N1+N2, ...), cu=cumulative + # max_len_q: scalar, max(N1, N2, ...) + q, indices_q, cu_seqlens_q, max_len_q = unpad_input(q, mask_q) + k, indices_kv, cu_seqlens_kv, max_len_kv = unpad_input(k, mask_kv) + v = index_first_axis(v.reshape(-1, H, D), indices_kv) # same indice as k + + # call varlen_func + out = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_len_q, + max_seqlen_k=max_len_kv, + dropout_p=dropout, + causal=causal, + window_size=window_size, + ) + + # pad (put back) output + out = pad_input(out, indices_q, B, N) + return out + else: + raise NotImplementedError("masked attention requires flash_attn!") + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x): + rnorm = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return (x * rnorm).to(dtype=self.weight.dtype) * self.weight + + +class SelfAttention(nn.Module): + def __init__( + self, + hidden_dim, + num_heads, + input_dim=None, + output_dim=None, + dropout=0, + causal=False, + qknorm=False, + qknorm_type="LayerNorm", + ): + super().__init__() + self.hidden_dim = hidden_dim + self.input_dim = input_dim if input_dim is not None else hidden_dim + self.output_dim = output_dim if output_dim is not None else hidden_dim + self.num_heads = num_heads + assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" + self.head_dim = hidden_dim // num_heads + self.causal = causal + self.dropout = dropout + self.qknorm = qknorm + + self.qkv_proj = nn.Linear(self.input_dim, 3 * self.hidden_dim) + self.out_proj = nn.Linear(self.hidden_dim, self.output_dim) + + if self.qknorm: + if qknorm_type == "RMSNorm": + self.q_norm = RMSNorm(self.hidden_dim, eps=1e-6) + self.k_norm = RMSNorm(self.hidden_dim, eps=1e-6) + else: + self.q_norm = nn.LayerNorm(self.hidden_dim, eps=1e-6, elementwise_affine=False) + self.k_norm = nn.LayerNorm(self.hidden_dim, eps=1e-6, elementwise_affine=False) + + def forward(self, x, mask=None): + # x: [B, N, C] + # mask: [B, N] + B, N, C = x.shape + qkv = self.qkv_proj(x) # [B, N, C] -> [B, N, 3 * D] + qkv = qkv.reshape(B, N, 3, -1).permute(2, 0, 1, 3) # [3, B, N, D] + q, k, v = qkv.chunk(3, dim=0) # [3, B, N, D] -> 3 * [1, B, N, D] + q = q.squeeze(0) + k = k.squeeze(0) + v = v.squeeze(0) + if self.qknorm: + q = self.q_norm(q) + k = self.k_norm(k) + q = q.reshape(B, N, self.num_heads, self.head_dim) + k = k.reshape(B, N, self.num_heads, self.head_dim) + v = v.reshape(B, N, self.num_heads, self.head_dim) + x = attention(q, k, v, mask_q=mask, mask_kv=mask, dropout=self.dropout, causal=self.causal) # [B, N, H, D] + x = self.out_proj(x.reshape(B, N, -1)) + return x + + +class CrossAttention(nn.Module): + def __init__( + self, + hidden_dim, + num_heads, + input_dim=None, + context_dim=None, + output_dim=None, + dropout=0, + qknorm=False, + qknorm_type="LayerNorm", + ): + super().__init__() + self.hidden_dim = hidden_dim + self.input_dim = input_dim if input_dim is not None else hidden_dim + self.context_dim = context_dim if context_dim is not None else hidden_dim + self.output_dim = output_dim if output_dim is not None else hidden_dim + self.num_heads = num_heads + assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" + self.head_dim = hidden_dim // num_heads + self.dropout = dropout + self.qknorm = qknorm + + self.q_proj = nn.Linear(self.input_dim, self.hidden_dim) + self.k_proj = nn.Linear(self.context_dim, self.hidden_dim) + self.v_proj = nn.Linear(self.context_dim, self.hidden_dim) + self.out_proj = nn.Linear(self.hidden_dim, self.output_dim) + + if self.qknorm: + if qknorm_type == "RMSNorm": + self.q_norm = RMSNorm(self.hidden_dim, eps=1e-6) + self.k_norm = RMSNorm(self.hidden_dim, eps=1e-6) + else: + self.q_norm = nn.LayerNorm(self.hidden_dim, eps=1e-6, elementwise_affine=False) + self.k_norm = nn.LayerNorm(self.hidden_dim, eps=1e-6, elementwise_affine=False) + + def forward(self, x, context, mask_q=None, mask_kv=None): + # x: [B, N, C] + # context: [B, M, C'] + # mask_q: [B, N] + # mask_kv: [B, M] + B, N, C = x.shape + M = context.shape[1] + q = self.q_proj(x) + k = self.k_proj(context) + v = self.v_proj(context) + if self.qknorm: + q = self.q_norm(q) + k = self.k_norm(k) + q = q.reshape(B, N, self.num_heads, self.head_dim) + k = k.reshape(B, M, self.num_heads, self.head_dim) + v = v.reshape(B, M, self.num_heads, self.head_dim) + x = attention(q, k, v, mask_q=mask_q, mask_kv=mask_kv, dropout=self.dropout, causal=False) # [B, N, H, D] + x = self.out_proj(x.reshape(B, N, -1)) + return x diff --git a/PartPacker/vae/modules/transformer.py b/PartPacker/vae/modules/transformer.py new file mode 100644 index 0000000..ac1c1f2 --- /dev/null +++ b/PartPacker/vae/modules/transformer.py @@ -0,0 +1,117 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +from .attention import CrossAttention, SelfAttention + + +class FeedForward(nn.Module): + def __init__(self, dim, mult=4): + super().__init__() + self.net = nn.Sequential(nn.Linear(dim, dim * mult), nn.GELU(), nn.Linear(dim * mult, dim)) + + def forward(self, x): + return self.net(x) + + +class AttentionBlock(nn.Module): + def __init__( + self, + dim, + num_heads, + dim_context=None, + qknorm=False, + gradient_checkpointing=True, + qknorm_type="LayerNorm", + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.dim_context = dim_context + self.gradient_checkpointing = gradient_checkpointing + + self.norm_attn = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + if dim_context is not None: + self.norm_context = nn.LayerNorm(dim_context, eps=1e-6, elementwise_affine=False) + self.attn = CrossAttention(dim, num_heads, context_dim=dim_context, qknorm=qknorm, qknorm_type=qknorm_type) + else: + self.attn = SelfAttention(dim, num_heads, qknorm=qknorm, qknorm_type=qknorm_type) + + self.norm_ff = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.ff = FeedForward(dim) + + def forward(self, x, c=None, mask=None, mask_c=None): + if self.training and self.gradient_checkpointing: + return checkpoint(self._forward, x, c, mask, mask_c, use_reentrant=False) + else: + return self._forward(x, c, mask, mask_c) + + def _forward(self, x, c=None, mask=None, mask_c=None): + # x: [B, N, C], hidden states + # c: [B, M, C'], condition (assume normed and projected to C) + # mask: [B, N], mask for x + # mask_c: [B, M], mask for c + # return: [B, N, C], updated hidden states + + if c is not None: + x = x + self.attn(self.norm_attn(x), self.norm_context(c), mask_q=mask, mask_kv=mask_c) + else: + x = x + self.attn(self.norm_attn(x), mask=mask) + + x = x + self.ff(self.norm_ff(x)) + + return x + + +# special attention block for the last cross-attn query layer +# 1. simple feed-forward (mult=1, no post ln) +# 2. no residual connection +# 3. no context ln +class FlashQueryLayer(nn.Module): + def __init__( + self, + dim, + num_heads, + dim_context, + qknorm=False, + gradient_checkpointing=True, + qknorm_type="LayerNorm", + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.dim_context = dim_context + self.gradient_checkpointing = gradient_checkpointing + + self.norm_attn = nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) + self.attn = CrossAttention(dim, num_heads, context_dim=dim_context, qknorm=qknorm, qknorm_type=qknorm_type) + self.ff = FeedForward(dim, mult=1) + + def forward(self, x, c=None, mask=None, mask_c=None): + if self.training and self.gradient_checkpointing: + return checkpoint(self._forward, x, c, mask, mask_c, use_reentrant=False) + else: + return self._forward(x, c, mask, mask_c) + + def _forward(self, x, c, mask=None, mask_c=None): + # x: [B, N, C], hidden states + # c: [B, M, C'], condition (assume normed and projected to C) + # mask: [B, N], mask for x + # mask_c: [B, M], mask for c + # return: [B, N, C], updated hidden states + + x = self.attn(self.norm_attn(x), c, mask_q=mask, mask_kv=mask_c) + x = self.ff(x) + + return x diff --git a/PartPacker/vae/scripts/infer.py b/PartPacker/vae/scripts/infer.py new file mode 100644 index 0000000..162c041 --- /dev/null +++ b/PartPacker/vae/scripts/infer.py @@ -0,0 +1,142 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import argparse +import glob +import importlib +import os +from datetime import datetime + +import fpsample +import kiui +import meshiki +import numpy as np +import torch +import trimesh + +from ..model import Model +from ..utils import box_normalize, postprocess_mesh, sphere_normalize, sync_timer + +# PYTHONPATH=. python vae/scripts/infer.py +parser = argparse.ArgumentParser() +parser.add_argument("--config", type=str, help="config file path", default="vae.configs.part_woenc") +parser.add_argument( + "--ckpt_path", + type=str, + help="checkpoint path", + default="pretrained/vae.pt", +) +parser.add_argument("--input", type=str, help="input directory", default="assets/meshes/") +parser.add_argument("--output_dir", type=str, help="output directory", default="output/") +parser.add_argument("--limit", type=int, help="how many samples to test", default=-1) +parser.add_argument("--num_fps_point", type=int, help="number of fps points", default=1024) +parser.add_argument("--num_fps_salient_point", type=int, help="number of fps salient points", default=1024) +parser.add_argument("--grid_res", type=int, help="grid resolution", default=512) +parser.add_argument("--seed", type=int, help="seed", default=42) +args = parser.parse_args() + + +TRIMESH_GLB_EXPORT = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).astype(np.float32) + +kiui.seed_everything(args.seed) + + +@sync_timer("prepare_input_from_mesh") +def prepare_input_from_mesh(mesh_path, use_salient_point=True, num_fps_point=1024, num_fps_salient_point=1024): + # load mesh, assume it's already processed to be watertight. + + mesh_name = mesh_path.split("/")[-1].split(".")[0] + vertices, faces = meshiki.load_mesh(mesh_path) + + # vertices = sphere_normalize(vertices) + vertices = box_normalize(vertices) + + mesh = meshiki.Mesh(vertices, faces) + + uniform_surface_points = mesh.uniform_point_sample(200000) + uniform_surface_points = meshiki.fps(uniform_surface_points, 32768) # hardcoded... + salient_surface_points = mesh.salient_point_sample(16384, thresh_bihedral=15) + + # save points + # trimesh.PointCloud(vertices=uniform_surface_points).export(os.path.join(workspace, mesh_name + "_uniform.ply")) + # trimesh.PointCloud(vertices=salient_surface_points).export(os.path.join(workspace, mesh_name + "_salient.ply")) + + sample = {} + + sample["pointcloud"] = torch.from_numpy(uniform_surface_points) + + # fps subsample + fps_indices = fpsample.bucket_fps_kdline_sampling(uniform_surface_points, num_fps_point, h=5, start_idx=0) + sample["fps_indices"] = torch.from_numpy(fps_indices).long() # [num_fps_point,] + + if use_salient_point: + sample["pointcloud_dorases"] = torch.from_numpy(salient_surface_points) # [N', 3] + + # fps subsample + fps_indices_dorases = fpsample.bucket_fps_kdline_sampling( + salient_surface_points, num_fps_salient_point, h=5, start_idx=0 + ) + sample["fps_indices_dorases"] = torch.from_numpy(fps_indices_dorases).long() # [num_fps_point,] + + return sample + + +print(f"Loading checkpoint from {args.ckpt_path}") +ckpt_dict = torch.load(args.ckpt_path, weights_only=True) + +# delete all keys other than model +if "model" in ckpt_dict: + ckpt_dict = ckpt_dict["model"] + +# instantiate model +print(f"Instantiating model from {args.config}") +model_config = importlib.import_module(args.config).make_config() +model = Model(model_config).eval().cuda().bfloat16() + +# load weight +print(f"Loading weights from {args.ckpt_path}") +model.load_state_dict(ckpt_dict, strict=True) + +timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") +workspace = os.path.join(args.output_dir, "vae_" + args.config.split(".")[-1] + "_" + timestamp) +if not os.path.exists(workspace): + os.makedirs(workspace) +else: + os.system(f"rm {workspace}/*") +print(f"Output directory: {workspace}") + +# load dataset +mesh_list = glob.glob(os.path.join(args.input, "*")) +mesh_list = mesh_list[: args.limit] if args.limit > 0 else mesh_list + +for i, mesh_path in enumerate(mesh_list): + print(f"Processing {i}/{len(mesh_list)}: {mesh_path}") + + mesh_name = mesh_path.split("/")[-1].split(".")[0] + + sample = prepare_input_from_mesh( + mesh_path, num_fps_point=args.num_fps_point, num_fps_salient_point=args.num_fps_salient_point + ) + for k in sample: + sample[k] = sample[k].unsqueeze(0).cuda() + + # call vae + with torch.inference_mode(): + output = model(sample, resolution=args.grid_res) + + latent = output["latent"] + vertices, faces = output["meshes"][0] + + mesh = trimesh.Trimesh(vertices, faces) + mesh = postprocess_mesh(mesh, 5e5) + + mesh.export(f"{workspace}/{mesh_name}.glb") diff --git a/PartPacker/vae/utils.py b/PartPacker/vae/utils.py new file mode 100644 index 0000000..068a974 --- /dev/null +++ b/PartPacker/vae/utils.py @@ -0,0 +1,315 @@ +""" +----------------------------------------------------------------------------- +Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +""" + +import os +from functools import wraps +from typing import Literal + +import numpy as np +import torch +import trimesh +from kiui.mesh_utils import clean_mesh, decimate_mesh + + +# Adapted from https://github.com/Tencent/Hunyuan3D-2/blob/main/hy3dgen/shapegen/utils.py#L38 +class sync_timer: + """ + Synchronized timer to count the inference time of `nn.Module.forward` or else. + set env var TIMER=1 to enable logging! + + Example as context manager: + ```python + with timer('name'): + run() + ``` + + Example as decorator: + ```python + @timer('name') + def run(): + pass + ``` + """ + + def __init__(self, name=None, flag_env="TIMER"): + self.name = name + self.flag_env = flag_env + + def __enter__(self): + if os.environ.get(self.flag_env, "0") == "1": + self.start = torch.cuda.Event(enable_timing=True) + self.end = torch.cuda.Event(enable_timing=True) + self.start.record() + return lambda: self.time + + def __exit__(self, exc_type, exc_value, exc_tb): + if os.environ.get(self.flag_env, "0") == "1": + self.end.record() + torch.cuda.synchronize() + self.time = self.start.elapsed_time(self.end) + if self.name is not None: + print(f"{self.name} takes {self.time} ms") + + def __call__(self, func): + @wraps(func) + def wrapper(*args, **kwargs): + with self: + result = func(*args, **kwargs) + return result + + return wrapper + + +@torch.no_grad() +def calculate_iou(pred: torch.Tensor, gt: torch.Tensor, target_value: int, thresh: float = 0) -> torch.Tensor: + """Calculate the Intersection over Union (IoU) between two volumes. + + Args: + pred (torch.Tensor): [*] continuous value between 0 and 1 + gt (torch.Tensor): [*] discrete value of 0 or 1 + target_value (int): The value to be considered as the target class + + Returns: + torch.Tensor: IoU value + """ + # Ensure volumes have the same shape + assert pred.shape == gt.shape, "Volumes must have the same shape" + + # binarize + pred_binary = pred > thresh + gt = gt > thresh + + # Convert the volumes to boolean tensors for logical operations + intersection = torch.logical_and(pred_binary == target_value, gt == target_value).sum().float() + union = torch.logical_or(pred_binary == target_value, gt == target_value).sum().float() + + # Compute IoU + iou = intersection / union if union != 0 else torch.tensor(0.0) + return iou + + +@torch.no_grad() +def calculate_metrics( + pred: torch.Tensor, gt: torch.Tensor, target_value: int = 1, thresh: float = 0.5 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Calculate Precision, Recall, and F1 between two volumes. + + Args: + pred (torch.Tensor): [*] continuous value between 0 and 1 + gt (torch.Tensor): [*] discrete value of 0 or 1 + target_value (int): The value to be considered as the target class + + Returns: + tuple: Precision, Recall, F1 values + """ + assert pred.shape == gt.shape, f"Pred {pred.shape} and gt {gt.shape} must have the same shape" + + # Binarize prediction + pred_binary = pred > thresh + gt = gt > thresh + + # True Positive (TP): pred == target_value and gt == target_value + true_positive = torch.logical_and(pred_binary == target_value, gt == target_value).sum().float() + + # False Positive (FP): pred == target_value and gt != target_value + false_positive = torch.logical_and(pred_binary == target_value, gt != target_value).sum().float() + + # False Negative (FN): pred != target_value and gt == target_value + false_negative = torch.logical_and(pred_binary != target_value, gt == target_value).sum().float() + + # Precision: TP / (TP + FP), best to detect False Positives + precision = ( + true_positive / (true_positive + false_positive) if (true_positive + false_positive) != 0 else torch.tensor(0.0) + ) + + # Recall: TP / (TP + FN), best to detect False Negatives + recall = ( + true_positive / (true_positive + false_negative) if (true_positive + false_negative) != 0 else torch.tensor(0.0) + ) + + # f1: 2 / (1 / precision + 1 / recall) + f1 = 2 / (1 / precision + 1 / recall) if (precision != 0 and recall != 0) else torch.tensor(0.0) + + return precision, recall, f1 + + +# Adapted from https://github.com/Stability-AI/stablediffusion/blob/main/ldm/modules/distributions/distributions.py#L24 +class DiagonalGaussianDistribution: + """VAE latent""" + + def __init__(self, mean, logvar, deterministic=False): + # mean, logvar: [B, L, D] x 2 + self.mean, self.logvar = mean, logvar + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean, device=self.mean.device, dtype=self.mean.dtype) + + def sample(self, weight: float = 1.0): + sample = weight * torch.randn(self.mean.shape, device=self.mean.device, dtype=self.mean.dtype) + x = self.mean + self.std * sample + return x + + def kl(self, other=None, dims=[1, 2]): + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.mean(torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=dims) + else: + return 0.5 * torch.mean( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=dims, + ) + + def nll(self, sample, dims=[1, 2]): + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.mean(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims) + + def mode(self): + return self.mean + + +class DummyLatent: + def __init__(self, mean): + self.mean = mean + + def sample(self, weight=0): + # simply perturb the mean + if weight > 0: + noise = torch.randn_like(self.mean) * weight + else: + noise = 0 + return self.mean + noise + + def mode(self): + return self.mean + + def kl(self): + # just an l2 penalty + return 0.5 * torch.mean(torch.pow(self.mean, 2)) + + +def construct_grid_points( + resolution: int, + indexing: str = "ij", +): + """Generate dense grid points in [-1, 1]^3. + + Args: + resolution (int): resolution of the grid + indexing (str, optional): indexing of the grid. Defaults to "ij". + + Returns: + torch.Tensor: grid points (resolution + 1, resolution + 1, resolution + 1, 3), inside bbox. + """ + x = np.linspace(-1, 1, resolution + 1, dtype=np.float32) + y = np.linspace(-1, 1, resolution + 1, dtype=np.float32) + z = np.linspace(-1, 1, resolution + 1, dtype=np.float32) + [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) + xyzs = np.stack((xs, ys, zs), axis=-1) + xyzs = torch.from_numpy(xyzs).float() + return xyzs + + +_diso_session = None # lazy session for reuse + + +@sync_timer("extract_mesh") +def extract_mesh( + grid_vals: torch.Tensor, + resolution: int, + isosurface_level: float = 0, + backend: Literal["mcubes", "diso"] = "mcubes", +): + """Extract mesh from grid occupancy. + + Args: + grid_vals (torch.Tensor): [resolution + 1, resolution + 1, resolution + 1], assume to be TSDF in [-1, 1] (inner is positive) + resolution (int, optional): Grid resolution. + isosurface_level (float, optional): Iso-surface level. Defaults to 0. + backend (Literal["mcubes", "diso"], optional): Backend for mesh extraction. Defaults to "diso", which uses GPU and is faster. + Returns: + vertices (np.ndarray): [N, 3], float32, in [-1, 1] + faces (np.ndarray): [M, 3], int32 + """ + + grid_vals = grid_vals.view(resolution + 1, resolution + 1, resolution + 1) + + if backend == "mcubes": + try: + import mcubes + except ImportError: + os.system("pip install pymcubes") + import mcubes + grid_vals = grid_vals.float().cpu().numpy() + verts, faces = mcubes.marching_cubes(grid_vals, isosurface_level) + verts = 2 * verts / resolution - 1.0 # normalize to [-1, 1] + elif backend == "diso": + try: + import diso + except ImportError: + os.system("pip install diso") + import diso + global _diso_session + if _diso_session is None: + _diso_session = diso.DiffDMC(dtype=torch.float32).cuda() + + grid_vals = -grid_vals.float().cuda() # diso assumes inner is NEGATIVE! + verts, faces = _diso_session(grid_vals, deform=None, normalize=True) # verts in [0, 1] + verts = verts.cpu().numpy() * 2 - 1.0 # normalize to [-1, 1] + faces = faces.cpu().numpy() + + return verts, faces + + +@sync_timer("postprocess_mesh") +def postprocess_mesh(mesh: trimesh.Trimesh, decimate_target=100000): + vertices = mesh.vertices + triangles = mesh.faces + + if vertices.shape[0] > 0 and triangles.shape[0] > 0: + vertices, triangles = clean_mesh(vertices, triangles, remesh=False, min_f=25, min_d=5) + if decimate_target > 0 and triangles.shape[0] > decimate_target: + vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False) + if vertices.shape[0] > 0 and triangles.shape[0] > 0: + vertices, triangles = clean_mesh(vertices, triangles, remesh=False, min_f=25, min_d=5) + + mesh.vertices = vertices + mesh.faces = triangles + + return mesh + + +def sphere_normalize(vertices): + bmin = vertices.min(axis=0) + bmax = vertices.max(axis=0) + bcenter = (bmax + bmin) / 2 + radius = np.linalg.norm(vertices - bcenter, axis=-1).max() + vertices = (vertices - bcenter) / radius # to [-1, 1] + return vertices + + +def box_normalize(vertices, bound=0.95): + bmin = vertices.min(axis=0) + bmax = vertices.max(axis=0) + bcenter = (bmax + bmin) / 2 + vertices = bound * (vertices - bcenter) / (bmax - bmin).max() + return vertices diff --git a/PartPacker_node.py b/PartPacker_node.py new file mode 100644 index 0000000..a60b743 --- /dev/null +++ b/PartPacker_node.py @@ -0,0 +1,145 @@ +# !/usr/bin/env python +# -*- coding: UTF-8 -*- +import os +import torch +import rembg +import numpy as np +from pathlib import PureWindowsPath +from .PartPacker.flow.configs.schema import ModelConfig +from .PartPacker.flow.model import Model +from .PartPacker.app import process_3d +from .node_utils import tensor2cv,gc_clear,add_mask,tensor2pil_upscale + + +import folder_paths + +MAX_SEED = np.iinfo(np.int32).max +current_node_path = os.path.dirname(os.path.abspath(__file__)) + + +device = torch.device( + "cuda:0") if torch.cuda.is_available() else torch.device( + "mps") if torch.backends.mps.is_available() else torch.device( + "cpu") + +# add checkpoints dir +PartPacker_Weigths_Path = os.path.join(folder_paths.models_dir, "PartPacker") +if not os.path.exists(PartPacker_Weigths_Path): + os.makedirs(PartPacker_Weigths_Path) +folder_paths.add_model_folder_path("PartPacker", PartPacker_Weigths_Path) + + +class PartPacker_Loader: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "checkpoint": (["none"] + [i for i in folder_paths.get_filename_list("PartPacker") if i.endswith(".pt")],), + "vae": (folder_paths.get_filename_list("vae"),), + "dino":("STRING", { "default": "facebook/dinov2-giant"}), + "cpu_offload": ("BOOLEAN", {"default": True},), + }, + } + + RETURN_TYPES = ("PartPacker_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "loader_main" + CATEGORY = "PartPacker" + + def loader_main(self, checkpoint,vae,dino,cpu_offload,): + if checkpoint == "none": + raise ValueError("No checkpoint selected") + + flow_ckpt_path=folder_paths.get_full_path("PartPacker", checkpoint) + vae_ckpt_path=folder_paths.get_full_path("vae", vae) + + # load model + print("***********Load model ***********") + TRIMESH_GLB_EXPORT = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).astype(np.float32) + bg_remover = rembg.new_session() + if not dino: + raise ValueError("No dino model path fill") + else: + if dino.count('/')!=1: + dino = PureWindowsPath(dino).as_posix() + dino_model="dinov2_vitg14" if "giant" in dino.lower() else "dinov2_vitl14_reg" + + # model config + model_config = ModelConfig( + vae_conf="custom_nodes.ComfyUI_PartPacker.PartPacker.vae.configs.part_woenc", + vae_ckpt_path=vae_ckpt_path, + qknorm=True, + qknorm_type="RMSNorm", + use_pos_embed=False, + dino_model=dino_model, + hidden_dim=1536, + flow_shift=3.0, + logitnorm_mean=1.0, + logitnorm_std=1.0, + latent_size=4096, + use_parts=True, + ) + + # instantiate model + model = Model(model_config,device,dino,cpu_offload=cpu_offload).eval().cuda().bfloat16() + + # load weight + ckpt_dict = torch.load(flow_ckpt_path, weights_only=True) + model.load_state_dict(ckpt_dict, strict=True) + + print("***********Load model done ***********") + gc_clear() + return ({"pipe": model, "bg_remover": bg_remover, "TRIMESH_GLB_EXPORT": TRIMESH_GLB_EXPORT},) + + + +class PartPacker_Sampler: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("PartPacker_MODEL",), + "image": ("IMAGE",),# BHWC + "target_num_faces" : ("INT", {"default": 100000, "min": 1000, "max": MAX_SEED, "step": 1}), + "grid_res": ("INT", {"default": 384, "min": 128, "max": 2048, "step": 16}), + "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}), + "steps": ("INT", {"default": 50, "min": 3, "max": 1024, "step": 1}), + "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1, "max": 20, "step": 0.1}), + "simplify_mesh": ("BOOLEAN", {"default": False},),}, + "optional":{ "mask": ("MASK",), # B H W 默认的mask是遮罩区黑色，而非传统的白色 + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("model_path", ) + FUNCTION = "sampler_main" + CATEGORY = "PartPacker" + + def sampler_main(self, model,image,target_num_faces,grid_res,seed,steps,cfg_scale,simplify_mesh,**kwargs): + if isinstance(kwargs.get("mask"),torch.Tensor): + mask=add_mask(kwargs.get("mask"),image) + mask=tensor2cv(mask) + else: + mask=None + input_image=tensor2cv(image,RGB2BGR=False) + model_path=process_3d(model.get("pipe"),model.get("bg_remover"),input_image,model.get("TRIMESH_GLB_EXPORT"), mask,folder_paths.get_output_directory(),num_steps=steps, cfg_scale=cfg_scale, grid_res=grid_res, seed=seed, simplify_mesh=simplify_mesh, target_num_faces=target_num_faces) + gc_clear() + return (model_path,) + + + +NODE_CLASS_MAPPINGS = { + "PartPacker_Loader": PartPacker_Loader, + "PartPacker_Sampler": PartPacker_Sampler, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PartPacker_Loader": "PartPacker_Loader", + "PartPacker_Sampler": "PartPacker_Sampler", +} diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..64146cd --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ + +from .PartPacker_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/example_workflows/ComfyUI_PartPacker.json b/example_workflows/ComfyUI_PartPacker_Workflow .json similarity index 57% rename from example_workflows/ComfyUI_PartPacker.json rename to example_workflows/ComfyUI_PartPacker_Workflow .json index a8de89a..730f47b 100644 --- a/example_workflows/ComfyUI_PartPacker.json +++ b/example_workflows/ComfyUI_PartPacker_Workflow .json @@ -1,130 +1,9 @@ { "id": "0ba45310-0a4e-4db4-869d-1b294e479e5b", "revision": 0, - "last_node_id": 9, - "last_link_id": 5, + "last_node_id": 17, + "last_link_id": 29, "nodes": [ - { - "id": 1, - "type": "PartPacker_Loader", - "pos": [ - 20743.24609375, - -1486.281982421875 - ], - "size": [ - 270, - 130 - ], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "model", - "type": "PartPacker_MODEL", - "links": [ - 1 - ] - } - ], - "properties": { - "Node name for S&R": "PartPacker_Loader" - }, - "widgets_values": [ - "flow.pt", - "vae.pt", - "F:\\test\\ComfyUI\\models\\diffusers\\facebook\\dinov2-giant", - true - ] - }, - { - "id": 7, - "type": "LoadImage", - "pos": [ - 20739.16015625, - -1297.8453369140625 - ], - "size": [ - 270, - 314 - ], - "flags": {}, - "order": 1, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 4 - ] - }, - { - "name": "MASK", - "type": "MASK", - "links": null - } - ], - "properties": { - "Node name for S&R": "LoadImage" - }, - "widgets_values": [ - "rabbit.png", - "image" - ] - }, - { - "id": 2, - "type": "PartPacker_Sampler", - "pos": [ - 21137.203125, - -1410.3333740234375 - ], - "size": [ - 286, - 288 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "model", - "type": "PartPacker_MODEL", - "link": 1 - }, - { - "name": "image", - "type": "IMAGE", - "link": 4 - } - ], - "outputs": [ - { - "name": "model_path", - "type": "STRING", - "links": [ - 5 - ] - } - ], - "properties": { - "Node name for S&R": "PartPacker_Sampler" - }, - "widgets_values": [ - 512, - 512, - 100000, - 384, - 788606917, - "randomize", - 50, - 7, - false - ] - }, { "id": 9, "type": "Preview3DAnimation", @@ -137,7 +16,7 @@ 550 ], "flags": {}, - "order": 3, + "order": 5, "mode": 0, "inputs": [ { @@ -152,18 +31,18 @@ "widget": { "name": "model_file" }, - "link": 5 + "link": 10 } ], "outputs": [], "properties": { "Node name for S&R": "Preview3DAnimation", - "Last Time Model File": "F:/ComfyUI311/ComfyUI/output/partpacker_20250617_180914.glb", + "Last Time Model File": "F:/ComfyUI311/ComfyUI/output/partpacker_20250617_221152.glb", "Camera Info": { "position": { - "x": -1.6016784775866804, - "y": 4.191211347281783, - "z": 14.285979669021703 + "x": 0.6679302476366114, + "y": 2.182846789244992, + "z": 13.06217089619584 }, "target": { "x": 0, @@ -175,45 +54,258 @@ } }, "widgets_values": [ - "F:/ComfyUI311/ComfyUI/output/partpacker_20250617_180914.glb", + "F:/ComfyUI311/ComfyUI/output/partpacker_20250617_221152.glb", "" ] + }, + { + "id": 7, + "type": "LoadImage", + "pos": [ + 20638.255859375, + -1336.9361572265625 + ], + "size": [ + 270, + 314 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 8 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": [] + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "rabbit.png", + "image" + ] + }, + { + "id": 1, + "type": "PartPacker_Loader", + "pos": [ + 20724.15234375, + -1517.1907958984375 + ], + "size": [ + 270, + 130 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "model", + "type": "PartPacker_MODEL", + "links": [ + 7 + ] + } + ], + "properties": { + "Node name for S&R": "PartPacker_Loader" + }, + "widgets_values": [ + "flow.pt", + "vae.pt", + "F:\\test\\ComfyUI\\models\\diffusers\\facebook\\dinov2-giant", + true + ] + }, + { + "id": 14, + "type": "LoadImageMask", + "pos": [ + 20637.064453125, + -951.5489501953125 + ], + "size": [ + 270, + 318 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 18 + ] + } + ], + "properties": { + "Node name for S&R": "LoadImageMask" + }, + "widgets_values": [ + "rabbit.png", + "alpha", + "image" + ] + }, + { + "id": 10, + "type": "PartPacker_Sampler", + "pos": [ + 21091.353515625, + -1372.093994140625 + ], + "size": [ + 350, + 252.90908813476562 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "PartPacker_MODEL", + "link": 7 + }, + { + "name": "image", + "type": "IMAGE", + "link": 8 + }, + { + "name": "mask", + "shape": 7, + "type": "MASK", + "link": null + } + ], + "outputs": [ + { + "name": "model_path", + "type": "STRING", + "links": [ + 10 + ] + } + ], + "properties": { + "Node name for S&R": "PartPacker_Sampler" + }, + "widgets_values": [ + 99854, + 384, + 568320646, + "randomize", + 50, + 7, + false + ] + }, + { + "id": 15, + "type": "InvertMask", + "pos": [ + 20935.978515625, + -1034.3560791015625 + ], + "size": [ + 140, + 26 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 18 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [] + } + ], + "properties": { + "Node name for S&R": "InvertMask" + }, + "widgets_values": [] } ], "links": [ [ - 1, + 7, 1, 0, - 2, + 10, 0, "PartPacker_MODEL" ], [ - 4, + 8, 7, 0, - 2, + 10, 1, "IMAGE" ], [ - 5, - 2, + 10, + 10, 0, 9, 1, "STRING" + ], + [ + 18, + 14, + 0, + 15, + 0, + "MASK" ] ], - "groups": [], + "groups": [ + { + "id": 1, + "title": "Group", + "bounding": [ + 20491.21484375, + -1613.211669921875, + 1379.64501953125, + 993.4290161132812 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], "config": {}, "extra": { "ds": { - "scale": 1.1, + "scale": 0.7513148009015777, "offset": [ - -20483.023344158355, - 1611.8175067860536 + -20162.192725743967, + 1672.318286488258 ] }, "frontendVersion": "1.21.6" diff --git a/example_workflows/__init__.py b/example_workflows/__init__.py deleted file mode 100644 index 8b13789..0000000 --- a/example_workflows/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/example_workflows/example.png b/example_workflows/example.png index 13f413e..bac5b81 100644 Binary files a/example_workflows/example.png and b/example_workflows/example.png differ diff --git a/example_workflows/example_new.png b/example_workflows/example_new.png deleted file mode 100644 index 985a07c..0000000 Binary files a/example_workflows/example_new.png and /dev/null differ diff --git a/node_utils.py b/node_utils.py new file mode 100644 index 0000000..e74bf05 --- /dev/null +++ b/node_utils.py @@ -0,0 +1,234 @@ +# !/usr/bin/env python +# -*- coding: UTF-8 -*- +import os +import torch +from PIL import Image +import numpy as np +import cv2 +import gc + +from comfy.utils import common_upscale,ProgressBar +from huggingface_hub import hf_hub_download + +cur_path = os.path.dirname(os.path.abspath(__file__)) +device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" + + +def gc_clear(): + gc.collect() + torch.cuda.empty_cache() + +def cv2pil(cv_image): + """ + 将OpenCV图像转换为PIL图像 + :param cv_image: OpenCV图像 + :return: PIL图像 + """ + # 将图像从BGR转换为RGB + rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) + # 使用PIL的Image.fromarray方法将NumPy数组转换为PIL图像 + pil_image = Image.fromarray(rgb_image) + return pil_image + + +def add_mask(masks,images): + if images.max() > 1.0: + white_value = 255.0 + else: + white_value = 1.0 + + if masks.max() > 1.0: + masks = masks / 255.0 + + masks = masks.unsqueeze(-1) + + white_bg = torch.ones_like(images) * white_value + + if images.shape[-1] == 4: + images = images[..., :3] + + tensor_image = images * masks + white_bg * (1 - masks) + return tensor_image + + + +def tensor_to_pil(tensor): + image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() + image = Image.fromarray(image_np, mode='RGB') + return image + +def tensor2pil_list(image,width,height): + B,_,_,_=image.size() + if B==1: + ref_image_list=[tensor2pil_upscale(image,width,height)] + else: + img_list = list(torch.chunk(image, chunks=B)) + ref_image_list = [tensor2pil_upscale(img,width,height) for img in img_list] + return ref_image_list + + +def tensor_upscale(img_tensor, width, height): + samples = img_tensor.movedim(-1, 1) + img = common_upscale(samples, width, height, "nearest-exact", "center") + samples = img.movedim(1, -1) + return samples + +def tensor2pil_upscale(img_tensor, width, height): + samples = img_tensor.movedim(-1, 1) + img = common_upscale(samples, width, height, "nearest-exact", "center") + samples = img.movedim(1, -1) + img_pil = tensor_to_pil(samples) + return img_pil + + +def tensor2cv(tensor_image,RGB2BGR=True): + if len(tensor_image.shape)==4:#bhwc to hwc + tensor_image=tensor_image.squeeze(0) + if tensor_image.is_cuda: + tensor_image = tensor_image.cpu().detach() + tensor_image=tensor_image.numpy() + #反归一化 + maxValue=tensor_image.max() + tensor_image=tensor_image*255/maxValue + img_cv2=np.uint8(tensor_image)#32 to uint8 + if RGB2BGR: + img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR) + return img_cv2 + +def cvargb2tensor(img): + assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) + img = torch.from_numpy(img.transpose((2, 0, 1))) + return img.float().div(255).unsqueeze(0) # 255也可以改为256 + +def cv2tensor(img): + assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = torch.from_numpy(img.transpose((2, 0, 1))) + return img.float().div(255).unsqueeze(0) # 255也可以改为256 + +def images_generator(img_list: list,): + #get img size + sizes = {} + for image_ in img_list: + if isinstance(image_,Image.Image): + count = sizes.get(image_.size, 0) + sizes[image_.size] = count + 1 + elif isinstance(image_,np.ndarray): + count = sizes.get(image_.shape[:2][::-1], 0) + sizes[image_.shape[:2][::-1]] = count + 1 + else: + raise "unsupport image list,must be pil or cv2!!!" + size = max(sizes.items(), key=lambda x: x[1])[0] + yield size[0], size[1] + + # any to tensor + def load_image(img_in): + if isinstance(img_in, Image.Image): + img_in=img_in.convert("RGB") + i = np.array(img_in, dtype=np.float32) + i = torch.from_numpy(i).div_(255) + if i.shape[0] != size[1] or i.shape[1] != size[0]: + i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0) + i = common_upscale(i, size[0], size[1], "lanczos", "center") + i = i.squeeze(0).movedim(0, -1).numpy() + return i + elif isinstance(img_in,np.ndarray): + i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32) + i = torch.from_numpy(i).div_(255) + #print(i.shape) + return i + else: + raise "unsupport image list,must be pil,cv2 or tensor!!!" + + total_images = len(img_list) + processed_images = 0 + pbar = ProgressBar(total_images) + images = map(load_image, img_list) + try: + prev_image = next(images) + while True: + next_image = next(images) + yield prev_image + processed_images += 1 + pbar.update_absolute(processed_images, total_images) + prev_image = next_image + except StopIteration: + pass + if prev_image is not None: + yield prev_image + +def load_images(img_list: list,): + gen = images_generator(img_list) + (width, height) = next(gen) + images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3))))) + if len(images) == 0: + raise FileNotFoundError(f"No images could be loaded .") + return images + +def tensor2pil(tensor): + image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() + image = Image.fromarray(image_np, mode='RGB') + return image + +def pil2narry(img): + narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) + return narry + +def equalize_lists(list1, list2): + """ + 比较两个列表的长度，如果不一致，则将较短的列表复制以匹配较长列表的长度。 + + 参数: + list1 (list): 第一个列表 + list2 (list): 第二个列表 + + 返回: + tuple: 包含两个长度相等的列表的元组 + """ + len1 = len(list1) + len2 = len(list2) + + if len1 == len2: + pass + elif len1 < len2: + print("list1 is shorter than list2, copying list1 to match list2's length.") + list1.extend(list1 * ((len2 // len1) + 1)) # 复制list1以匹配list2的长度 + list1 = list1[:len2] # 确保长度一致 + else: + print("list2 is shorter than list1, copying list2 to match list1's length.") + list2.extend(list2 * ((len1 // len2) + 1)) # 复制list2以匹配list1的长度 + list2 = list2[:len1] # 确保长度一致 + + return list1, list2 + +def file_exists(directory, filename): + # 构建文件的完整路径 + file_path = os.path.join(directory, filename) + # 检查文件是否存在 + return os.path.isfile(file_path) + +def download_weights(file_dir,repo_id,subfolder="",pt_name=""): + if subfolder: + file_path = os.path.join(file_dir,subfolder, pt_name) + sub_dir=os.path.join(file_dir,subfolder) + if not os.path.exists(sub_dir): + os.makedirs(sub_dir) + if not os.path.exists(file_path): + file_path = hf_hub_download( + repo_id=repo_id, + subfolder=subfolder, + filename=pt_name, + local_dir = file_dir, + ) + return file_path + else: + file_path = os.path.join(file_dir, pt_name) + if not os.path.exists(file_dir): + os.makedirs(file_dir) + if not os.path.exists(file_path): + file_path = hf_hub_download( + repo_id=repo_id, + filename=pt_name, + local_dir=file_dir, + ) + return file_path diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f350311 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "ComfyUI_PartPacker" +description = "This is the comfyui implementation of PartPacker: Efficient Part-level 3D Object Generation via Dual Volume Packing.Max varm12G" +version = "1.0.0" +license = {file = "LICENSE"} +dependencies = ["torch", "numpy", "trimesh", "#fpsample", "einops", "onnxruntime", "rembg", "kiui", "torchvision", "pymcubes", "tqdm", "opencv-python", "pymeshlab", "#gradio", "transformers"] + +[project.urls] +Repository = "https://github.com/smthemex/ComfyUI_PartPacker" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "smthemex" +DisplayName = "ComfyUI" +Icon = ""