add partcrafter scene

This commit is contained in:
Denis Kochetov
2025-07-29 18:34:02 +00:00
parent ffa5ff828e
commit 06930669e2
9 changed files with 417 additions and 0 deletions
@@ -0,0 +1,35 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
@@ -0,0 +1,18 @@
---
license: mit
pipeline_tag: image-to-3d
tags:
- partcrafter
- 3d-generation
- rectified-flow
---
# PartCrafter: Structured 3D Mesh Generation via Compositional Latent Diffusion Transformers (Scene Level)
PartCrafter is a structured 3D generative model that jointly generates multiple parts and objects from a single RGB image in one shot.
This HuggingFace🤗 repo stores pretrained model weights for PartCrafter on **scene level**.
For more details about usage, please refer to the [official GitHub repo](https://github.com/wgsxm/PartCrafter).
- Project page: https://wgsxm.github.io/projects/partcrafter/
- Code: https://github.com/wgsxm/PartCrafter
- Paper: https://arxiv.org/abs/2506.05573
@@ -0,0 +1,27 @@
{
"crop_size": {
"height": 224,
"width": 224
},
"do_center_crop": true,
"do_convert_rgb": true,
"do_normalize": true,
"do_rescale": true,
"do_resize": true,
"image_mean": [
0.485,
0.456,
0.406
],
"image_processor_type": "BitImageProcessor",
"image_std": [
0.229,
0.224,
0.225
],
"resample": 3,
"rescale_factor": 0.00392156862745098,
"size": {
"shortest_edge": 256
}
}
@@ -0,0 +1,60 @@
{
"apply_layernorm": true,
"architectures": [
"Dinov2Model"
],
"attention_probs_dropout_prob": 0.0,
"drop_path_rate": 0.0,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.0,
"hidden_size": 1024,
"image_size": 518,
"initializer_range": 0.02,
"layer_norm_eps": 1e-06,
"layerscale_value": 1.0,
"mlp_ratio": 4,
"model_type": "dinov2",
"num_attention_heads": 16,
"num_channels": 3,
"num_hidden_layers": 24,
"out_features": [
"stage24"
],
"out_indices": [
24
],
"patch_size": 14,
"qkv_bias": true,
"reshape_hidden_states": true,
"stage_names": [
"stem",
"stage1",
"stage2",
"stage3",
"stage4",
"stage5",
"stage6",
"stage7",
"stage8",
"stage9",
"stage10",
"stage11",
"stage12",
"stage13",
"stage14",
"stage15",
"stage16",
"stage17",
"stage18",
"stage19",
"stage20",
"stage21",
"stage22",
"stage23",
"stage24"
],
"torch_dtype": "float16",
"transformers_version": "4.53.2",
"use_mask_token": true,
"use_swiglu_ffn": false
}
@@ -0,0 +1,24 @@
{
"_class_name": "PartCrafterPipeline",
"_diffusers_version": "0.34.0",
"feature_extractor_dinov2": [
"transformers",
"BitImageProcessor"
],
"image_encoder_dinov2": [
"transformers",
"Dinov2Model"
],
"scheduler": [
"partcrafter_src.schedulers.scheduling_rectified_flow",
"RectifiedFlowScheduler"
],
"transformer": [
"partcrafter_src.models.transformers.partcrafter_transformer",
"PartCrafterDiTModel"
],
"vae": [
"partcrafter_src.models.autoencoders.autoencoder_kl_triposg",
"TripoSGVAEModel"
]
}
@@ -0,0 +1,7 @@
{
"_class_name": "RectifiedFlowScheduler",
"_diffusers_version": "0.34.0",
"num_train_timesteps": 1000,
"shift": 1,
"use_dynamic_shifting": false
}
@@ -0,0 +1,27 @@
{
"_class_name": "PartCrafterDiTModel",
"_diffusers_version": "0.34.0",
"cross_attention_dim": 1024,
"enable_global_cross_attn": true,
"enable_local_cross_attn": true,
"enable_part_embedding": true,
"global_attn_block_id_range": null,
"global_attn_block_ids": [
0,
2,
4,
6,
8,
10,
12,
14,
16,
18,
20
],
"in_channels": 64,
"max_num_parts": 32,
"num_attention_heads": 16,
"num_layers": 21,
"width": 2048
}
@@ -0,0 +1,15 @@
{
"_class_name": "TripoSGVAEModel",
"_diffusers_version": "0.34.0",
"_name_or_path": "pretrained_weights/TripoSG",
"embed_frequency": 8,
"embed_include_pi": false,
"embedding_type": "frequency",
"in_channels": 3,
"latent_channels": 64,
"num_attention_heads": 8,
"num_layers_decoder": 16,
"num_layers_encoder": 8,
"width_decoder": 1024,
"width_encoder": 512
}
+204
View File
@@ -5790,4 +5790,208 @@ class PartCrafter_Generate:
print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}")
return (zip_path, relative_scene_path, processed_image_tensor)
#------ partcrafter scene ---------------------
class Load_PartCrafter_Scene_Pipeline:
"""Load PartCrafter Scene Pipeline"""
CATEGORY = "Comfy3D/Algorithm/PartCrafter"
RETURN_TYPES = ("DIFFUSERS_PIPE",)
RETURN_NAMES = ("partcrafter_scene_pipe",)
FUNCTION = "load"
_REPO_ID = "wgsxm/PartCrafter-Scene"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
}
@staticmethod
def _ensure_weights():
safe_repo_name = Load_PartCrafter_Scene_Pipeline._REPO_ID.replace("/", "_")
base_dir = os.path.join(CKPT_DIFFUSERS_PATH, Load_PartCrafter_Scene_Pipeline._REPO_ID)
required_files = [
"model_index.json",
"transformer/diffusion_pytorch_model.safetensors",
"vae/diffusion_pytorch_model.safetensors",
]
files_missing = []
for file_path in required_files:
full_file_path = os.path.join(base_dir, file_path)
if not os.path.exists(full_file_path):
files_missing.append(file_path)
if files_missing or not os.path.exists(base_dir):
print(f"Loading PartCrafter-Scene from {Load_PartCrafter_Scene_Pipeline._REPO_ID}...")
print(f"Missing files: {len(files_missing)}")
snapshot_download(
repo_id=Load_PartCrafter_Scene_Pipeline._REPO_ID,
repo_type="model",
local_dir=base_dir,
resume_download=True,
ignore_patterns=HF_DOWNLOAD_IGNORE,
)
print(f"PartCrafter-Scene loaded successfully")
else:
print(f"PartCrafter-Scene weights already loaded")
return base_dir
def load(self):
base_dir = self._ensure_weights()
pipeline = PartCrafterPipeline.from_pretrained(base_dir).to(DEVICE, WEIGHT_DTYPE)
print(f"PartCrafter-Scene pipeline loaded")
return (pipeline,)
class PartCrafter_Generate:
"""PartCrafter Generation - Creates multi-part 3D scenes with colored components"""
CATEGORY = "Comfy3D/Algorithm/PartCrafter"
RETURN_TYPES = ("STRING", "STRING", "IMAGE")
RETURN_NAMES = ("parts_zip_path", "glb_mesh_path", "processed_image")
FUNCTION = "generate"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"partcrafter_pipe": ("DIFFUSERS_PIPE",),
"image": ("IMAGE",),
"num_parts": ("INT", {"default": 4, "min": 1, "max": 16}),
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
"num_tokens": ("INT", {"default": 1024, "min": 256, "max": 2048}),
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 100}),
"guidance_scale": ("FLOAT", {"default": 7.0, "min": 0.0, "step": 0.1}),
"max_num_expanded_coords": ("INT", {"default": 1000000000, "min": 1000, "max": 10000000000}),
"use_flash_decoder": ("BOOLEAN", {"default": False}),
"remove_background": ("BOOLEAN", {"default": True}),
"sampling_version": ("INT", {"default": 1, "min": 1, "max": 2}),
}
}
@torch.no_grad()
def generate(self, partcrafter_pipe, image, num_parts, seed, num_tokens, num_inference_steps,
guidance_scale, max_num_expanded_coords, use_flash_decoder, remove_background, sampling_version):
# Convert image
pil_image = torch_imgs_to_pils(image)[0]
# Remove background if needed
if remove_background:
rmbg_worker = BackgroundRemover_2_1()
if pil_image.mode == "RGBA":
pil_image = pil_image.convert('RGB')
pil_image = rmbg_worker(pil_image)
del rmbg_worker
# Set sampling version
if hasattr(partcrafter_pipe.vae, 'set_sampling_version'):
partcrafter_pipe.vae.set_sampling_version(sampling_version)
# Generation
generator = torch.Generator(device=partcrafter_pipe.device)
generator = generator.manual_seed(int(seed))
outputs = partcrafter_pipe(
image=[pil_image] * num_parts,
attention_kwargs={"num_parts": num_parts},
num_tokens=num_tokens,
generator=generator,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
max_num_expanded_coords=max_num_expanded_coords,
use_flash_decoder=use_flash_decoder,
).meshes
# Ensure no None outputs
for i, mesh in enumerate(outputs):
if mesh is None:
outputs[i] = trimesh.Trimesh(vertices=[[0,0,0]], faces=[[0,0,0]])
print(f"Replaced None mesh at index {i} with dummy mesh")
merged_mesh_trimesh = get_colored_mesh_composition(outputs)
split_mesh = explode_mesh(merged_mesh_trimesh)
# Debug logging
print(f"PartCrafter result type: {type(merged_mesh_trimesh)}")
if hasattr(merged_mesh_trimesh, 'geometry'):
print(f"Scene contains {len(merged_mesh_trimesh.geometry)} parts: {list(merged_mesh_trimesh.geometry.keys())}")
# Create ZIP with individual parts
parts_output_dir = "output/partcrafter_parts"
os.makedirs(parts_output_dir, exist_ok=True)
zip_path = os.path.join(parts_output_dir, f"parts.zip")
parts = []
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
for idx, mesh in enumerate(outputs):
print(f"Part {idx} type: {type(mesh)}")
if mesh is not None:
# Save part as GLB file
part_path = os.path.join(parts_output_dir, f"part_{idx:02d}.glb")
mesh.export(part_path)
parts.append(part_path)
# Add to ZIP
zipf.write(part_path, f"part_{idx:02d}.glb")
# Clean up temporary file
try:
os.remove(part_path)
except:
pass
print(f"Created parts ZIP: {zip_path} with {len(parts)} parts")
# Save merged mesh
scene_output_dir = "output/partcrafter_scenes"
os.makedirs(scene_output_dir, exist_ok=True)
scene_file_path = os.path.join(scene_output_dir, f"scene.glb")
# For Preview_3DMesh, return relative path from ComfyUI output directory
relative_scene_path = "partcrafter_scenes/scene.glb"
# Export the merged mesh (same as app.py: merged.export(merged_path))
merged_mesh_trimesh.export(scene_file_path)
print(f"Saved merged colored mesh to: {scene_file_path}")
# Debug: check what was saved
if isinstance(merged_mesh_trimesh, trimesh.Scene):
print(f"Exported Scene with {len(merged_mesh_trimesh.geometry)} parts")
for geom_name in merged_mesh_trimesh.geometry:
geom = merged_mesh_trimesh.geometry[geom_name]
print(f" Part {geom_name}: vertices: {len(geom.vertices)}, faces: {len(geom.faces)}")
if hasattr(geom.visual, 'vertex_colors') and geom.visual.vertex_colors is not None:
print(f" Has vertex colors: {geom.visual.vertex_colors.shape}")
else:
print(f"Exported single mesh: vertices: {len(merged_mesh_trimesh.vertices)}, faces: {len(merged_mesh_trimesh.faces)}")
# Debug info about the scene
try:
print(f"Scene file saved successfully: {scene_file_path}")
if os.path.exists(scene_file_path):
file_size = os.path.getsize(scene_file_path) / (1024 * 1024) # MB
print(f"Scene file size: {file_size:.2f} MB")
print(f"Relative path for Preview_3DMesh: {relative_scene_path}")
else:
print(f"Warning: Scene file was not created properly")
except Exception as e:
print(f"Error checking scene file: {e}")
# Processed image
processed_image_tensor = pils_to_torch_imgs([pil_image])
print(f"PartCrafter: Generated {len(outputs)} parts with explode_mesh processing")
print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}")
return (zip_path, relative_scene_path, processed_image_tensor)