| """Convert the original Idol checkpoints (speridlabs/idol) into the diffusers weights layout used by these blocks: |
| `transformer/`, `octree_decoder/` and `gaussian_decoder/`, each with a `config.json` and safetensors. |
| |
| It imports the model classes next to it, so download it together with them: |
| |
| hf download OzzyGT/idol_custom_blocks convert_idol_to_diffusers.py transformer_idol.py octree_decoder_idol.py \ |
| gaussian_decoder_idol.py --local-dir idol_convert |
| cd idol_convert |
| python convert_idol_to_diffusers.py --output_dir ./idol_diffusers |
| python convert_idol_to_diffusers.py --checkpoint /path/to/speridlabs-idol --output_dir ./idol_diffusers |
| |
| Then point the pipeline at the result with |
| `pipe.load_components(pretrained_model_name_or_path={"transformer": out, "octree_decoder": out, "gaussian_decoder": out})` |
| or upload the folder and change `IDOL_WEIGHTS_REPO` in `components.py`. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import re |
|
|
| import torch |
| from gaussian_decoder_idol import IdolGaussianDecoderModel |
| from octree_decoder_idol import IdolOctreeDecoderModel |
| from safetensors.torch import load_file |
| from transformer_idol import IdolTransformerModel |
|
|
|
|
| ORIGINAL_FILES = ["config.json", "denoiser.safetensors", "decoder.safetensors", "decoder_gs.safetensors"] |
|
|
|
|
| def split_fused(state_dict, old_prefix, new_prefixes): |
| """Split a fused linear `old_prefix` along its output rows into the linears `new_prefixes`.""" |
| for suffix in ("weight", "bias"): |
| tensor = state_dict.pop(f"{old_prefix}.{suffix}") |
| for prefix, chunk in zip(new_prefixes, tensor.chunk(len(new_prefixes), dim=0)): |
| state_dict[f"{prefix}.{suffix}"] = chunk.contiguous() |
|
|
|
|
| def rename(state_dict, rules): |
| converted = {} |
| for key, value in state_dict.items(): |
| for pattern, replacement in rules: |
| key = re.sub(pattern, replacement, key) |
| converted[key] = value |
| return converted |
|
|
|
|
| def convert_attention(state_dict, old, new, fused_qkv): |
| if fused_qkv: |
| split_fused(state_dict, f"{old}.to_qkv", [f"{new}.to_q", f"{new}.to_k", f"{new}.to_v"]) |
| else: |
| state_dict[f"{new}.to_q.weight"] = state_dict.pop(f"{old}.to_q.weight") |
| state_dict[f"{new}.to_q.bias"] = state_dict.pop(f"{old}.to_q.bias") |
| split_fused(state_dict, f"{old}.to_kv", [f"{new}.to_k", f"{new}.to_v"]) |
| state_dict[f"{new}.to_out.0.weight"] = state_dict.pop(f"{old}.to_out.weight") |
| state_dict[f"{new}.to_out.0.bias"] = state_dict.pop(f"{old}.to_out.bias") |
| state_dict[f"{new}.norm_q.weight"] = state_dict.pop(f"{old}.q_rms_norm.gamma") |
| state_dict[f"{new}.norm_k.weight"] = state_dict.pop(f"{old}.k_rms_norm.gamma") |
|
|
|
|
| def num_blocks(state_dict): |
| return len({key.split(".")[1] for key in state_dict if key.startswith("blocks.")}) |
|
|
|
|
| def convert_transformer_state_dict(state_dict): |
| state_dict = rename( |
| state_dict, |
| [ |
| (r"^t_embedder\.mlp\.0\.", "t_embedder.linear_1."), |
| (r"^t_embedder\.mlp\.2\.", "t_embedder.linear_2."), |
| (r"^cam_refiner\.mlp\.", "cam_refiner.net."), |
| (r"\.mlp\.mlp\.", ".ff.net."), |
| (r"\.attn\.out\.", ".attn.to_out.0."), |
| (r"\.attn\.q_norm\.gamma$", ".attn.norm_q.weight"), |
| (r"\.attn\.k_norm\.gamma$", ".attn.norm_k.weight"), |
| (r"shift_table$", "scale_shift_table"), |
| ], |
| ) |
| for key in [k for k in state_dict if k.endswith(".attn.qkv.weight")]: |
| prefix = key[: -len(".qkv.weight")] |
| split_fused(state_dict, f"{prefix}.qkv", [f"{prefix}.to_q", f"{prefix}.to_k", f"{prefix}.to_v"]) |
| return state_dict |
|
|
|
|
| def convert_octree_decoder_state_dict(state_dict): |
| state_dict = rename( |
| state_dict, |
| [ |
| (r"^l_embedder\.mlp\.0\.", "l_embedder.linear_1."), |
| (r"^l_embedder\.mlp\.2\.", "l_embedder.linear_2."), |
| (r"\.mlp\.mlp\.", ".ff.net."), |
| ], |
| ) |
| for i in range(num_blocks(state_dict)): |
| convert_attention(state_dict, f"blocks.{i}.cross_attn", f"blocks.{i}.cross_attn", fused_qkv=False) |
| return state_dict |
|
|
|
|
| def convert_gaussian_decoder_state_dict(state_dict): |
| state_dict = rename(state_dict, [(r"\.mlp\.mlp\.", ".ff.net.")]) |
| for i in range(num_blocks(state_dict)): |
| convert_attention(state_dict, f"blocks.{i}.self_attn", f"blocks.{i}.attn1", fused_qkv=True) |
| convert_attention(state_dict, f"blocks.{i}.cross_attn", f"blocks.{i}.attn2", fused_qkv=False) |
| return state_dict |
|
|
|
|
| def transformer_config(original): |
| args = original["denoiser"]["args"] |
| return { |
| "num_latent_tokens": args["q_token_length"], |
| "in_channels": args["in_channels"], |
| "camera_channels": args["cam_channels"], |
| "hidden_size": args["model_channels"], |
| "image_embed_dim": args["cond_channels"], |
| "image_vae_embed_dim": args["cond2_channels"], |
| "num_layers": args["num_blocks"], |
| "num_refiner_layers": args["num_refiner_blocks"], |
| "num_attention_heads": args["num_heads"], |
| "attention_head_dim": args.get("num_head_channels", 64), |
| "mlp_ratio": float(args["mlp_ratio"]), |
| } |
|
|
|
|
| def octree_decoder_config(original): |
| args = original["decoder"]["args"] |
| return { |
| "hidden_size": args["model_channels"], |
| "latent_channels": args["cond_channels"], |
| "num_layers": args["num_blocks"], |
| "num_attention_heads": args["num_heads"], |
| "mlp_ratio": float(args["mlp_ratio"]), |
| "max_level": original["max_voxel_level"], |
| } |
|
|
|
|
| def gaussian_decoder_config(original): |
| args = original["decoder_gs"]["args"] |
| representation = args["representation_config"] |
| return { |
| "in_channels": args["in_channels"], |
| "hidden_size": args["model_channels"], |
| "latent_channels": args["cond_channels"], |
| "num_layers": args["num_blocks"], |
| "num_attention_heads": args["num_heads"], |
| "mlp_ratio": float(args["mlp_ratio"]), |
| "gaussians_per_point": representation["num_gaussians"], |
| "xyz_lr": representation["lr"]["_xyz"], |
| "features_dc_lr": representation["lr"]["_features_dc"], |
| "opacity_lr": representation["lr"]["_opacity"], |
| "scaling_lr": representation["lr"]["_scaling"], |
| "rotation_lr": representation["lr"]["_rotation"], |
| "perturb_size": representation["perturbe_size"], |
| "filter_kernel_size": representation["3d_filter_kernel_size"], |
| "scaling_bias": representation["scaling_bias"], |
| "opacity_bias": representation["opacity_bias"], |
| } |
|
|
|
|
| |
| COMPONENTS = { |
| "transformer": (IdolTransformerModel, "denoiser", convert_transformer_state_dict, transformer_config), |
| "octree_decoder": (IdolOctreeDecoderModel, "decoder", convert_octree_decoder_state_dict, octree_decoder_config), |
| "gaussian_decoder": ( |
| IdolGaussianDecoderModel, |
| "decoder_gs", |
| convert_gaussian_decoder_state_dict, |
| gaussian_decoder_config, |
| ), |
| } |
|
|
|
|
| def resolve_checkpoint(checkpoint): |
| if os.path.isdir(checkpoint): |
| return checkpoint |
| from huggingface_hub import snapshot_download |
|
|
| return snapshot_download(checkpoint, allow_patterns=ORIGINAL_FILES) |
|
|
|
|
| def convert_component(subfolder, checkpoint_dir): |
| """Build the diffusers model for `subfolder` from the original checkpoint (float32, on CPU).""" |
| model_cls, original_name, convert_state_dict, make_config = COMPONENTS[subfolder] |
| with open(os.path.join(checkpoint_dir, "config.json")) as f: |
| original_config = json.load(f) |
| model = model_cls(**make_config(original_config)) |
| original_state_dict = load_file(os.path.join(checkpoint_dir, f"{original_name}.safetensors")) |
| model.load_state_dict(convert_state_dict(original_state_dict), strict=True) |
| return model.float().eval() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--checkpoint", default="speridlabs/idol", help="Hub repo or local folder of the original") |
| parser.add_argument("--output_dir", required=True) |
| args = parser.parse_args() |
|
|
| checkpoint_dir = resolve_checkpoint(args.checkpoint) |
| for subfolder, (model_cls, *_) in COMPONENTS.items(): |
| model = convert_component(subfolder, checkpoint_dir) |
| out_dir = os.path.join(args.output_dir, subfolder) |
| model.save_pretrained(out_dir) |
|
|
| reloaded = model_cls.from_pretrained(out_dir) |
| reference = model.state_dict() |
| for key, value in reloaded.state_dict().items(): |
| assert torch.equal(value, reference[key]), f"{subfolder}: {key} changed after save/load" |
| num_params = sum(p.numel() for p in model.parameters()) / 1e6 |
| print(f"{subfolder}: {num_params:.1f}M parameters -> {out_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
