From 06930669e296ca7d8d9f8222db8d28b72f5fead4 Mon Sep 17 00:00:00 2001 From: Denis Kochetov Date: Tue, 29 Jul 2025 18:34:02 +0000 Subject: [PATCH] add partcrafter scene --- .../wgsxm/PartCrafter-Scene/.gitattributes | 35 +++ .../wgsxm/PartCrafter-Scene/README.md | 18 ++ .../preprocessor_config.json | 27 +++ .../image_encoder_dinov2/config.json | 60 ++++++ .../wgsxm/PartCrafter-Scene/model_index.json | 24 +++ .../scheduler/scheduler_config.json | 7 + .../PartCrafter-Scene/transformer/config.json | 27 +++ .../wgsxm/PartCrafter-Scene/vae/config.json | 15 ++ nodes.py | 204 ++++++++++++++++++ 9 files changed, 417 insertions(+) create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/.gitattributes create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/README.md create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/feature_extractor_dinov2/preprocessor_config.json create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/image_encoder_dinov2/config.json create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/model_index.json create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/scheduler/scheduler_config.json create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/transformer/config.json create mode 100644 Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/vae/config.json diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/.gitattributes b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/.gitattributes new file mode 100644 index 0000000..a6344aa --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/.gitattributes @@ -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 diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/README.md b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/README.md new file mode 100644 index 0000000..c76298b --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/README.md @@ -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 \ No newline at end of file diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/feature_extractor_dinov2/preprocessor_config.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/feature_extractor_dinov2/preprocessor_config.json new file mode 100644 index 0000000..ff5b47c --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/feature_extractor_dinov2/preprocessor_config.json @@ -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 + } +} diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/image_encoder_dinov2/config.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/image_encoder_dinov2/config.json new file mode 100644 index 0000000..72ae708 --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/image_encoder_dinov2/config.json @@ -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 +} diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/model_index.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/model_index.json new file mode 100644 index 0000000..1a846e7 --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/model_index.json @@ -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" + ] +} diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/scheduler/scheduler_config.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/scheduler/scheduler_config.json new file mode 100644 index 0000000..1114433 --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/scheduler/scheduler_config.json @@ -0,0 +1,7 @@ +{ + "_class_name": "RectifiedFlowScheduler", + "_diffusers_version": "0.34.0", + "num_train_timesteps": 1000, + "shift": 1, + "use_dynamic_shifting": false +} diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/transformer/config.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/transformer/config.json new file mode 100644 index 0000000..5aa5692 --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/transformer/config.json @@ -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 +} diff --git a/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/vae/config.json b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/vae/config.json new file mode 100644 index 0000000..4e372be --- /dev/null +++ b/Checkpoints/Diffusers/wgsxm/PartCrafter-Scene/vae/config.json @@ -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 +} diff --git a/nodes.py b/nodes.py index 2d1b55b..c5eddc9 100644 --- a/nodes.py +++ b/nodes.py @@ -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)