"""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"], } # subfolder -> (model class, original checkpoint name, state dict converter, config builder) 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()