#!/usr/bin/env python """ 70_convert_checkpoint.py — Convert a DeepSpeed pipeline-parallel checkpoint to an Anima-loadable safetensors file. Why this exists: diffusion-pipe's mid-training checkpoints are deepspeed pipeline-parallel format (one .pt file per pipeline "layer"). The clean .safetensors format only appears at epoch boundaries via the [save] phase, and our epoch boundary (step ~7550 for full 4-variant pass) is past our max_steps cap. To run a mid-training checkpoint in ComfyUI / inference for comparison purposes, we need to consolidate the layer .pt files into one .safetensors with the same key structure as the original Anima Preview3 base weights. The mapping was reverse-engineered by inspecting the structure: DeepSpeed file Contents Anima safetensors prefix --------------- -------- ------------------------ layer_00-model_states.pt x_embedder.*, net.* (no rename) t_embedder.*, t_embedding_norm.* layer_01-model_states.pt (empty) skip layer_02-model_states.pt block.self_attn.*, net.blocks.0.* (drop "block.") block.cross_attn.*, block.mlp.*, block.adaln_modulation_* layer_03-model_states.pt same shape net.blocks.1.* ... layer_29-model_states.pt same shape net.blocks.27.* layer_30-model_states.pt final_layer.* net.* (no rename) layer_31-model_states.pt (empty, if present) skip So the general rules: - layer_00 keys: prepend "net." - layer_NN where N in [2..N_blocks+1]: rename "block." → "net.blocks.." - last meaningful layer (containing final_layer.*): prepend "net." - empty layers: skip After conversion, the script verifies: 1. Output key set matches Anima base key set exactly (set equality) 2. Per-key tensor shapes match between converted output and Anima base Usage: python 70_convert_checkpoint.py \\ --deepspeed-dir /opt/local/outputs/anima_finish_v1//global_step1670 \\ --reference-safetensors /opt/local/models/anima/split_files/diffusion_models/anima-preview3-base.safetensors \\ --output /workspace/checkpoints/anima_finish_v1_step1670.safetensors Exit codes: 0 = converted + verified 1 = converted but verification failed 2 = no valid checkpoint files found 3 = ref file missing """ from __future__ import annotations import argparse import sys import time from collections import OrderedDict from pathlib import Path import torch from safetensors import safe_open from safetensors.torch import save_file def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) p.add_argument("--deepspeed-dir", required=True, type=Path, help="Directory containing layer_NN-model_states.pt files.") p.add_argument("--reference-safetensors", required=True, type=Path, help="Original Anima safetensors for key/shape verification.") p.add_argument("--output", required=True, type=Path, help="Path to write the converted safetensors.") p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"], help="Output tensor dtype. Default bfloat16 matches Anima.") p.add_argument("--skip-verify", action="store_true", help="Skip key/shape verification (faster, but unsafe).") p.add_argument("--merge-frozen-from-ref", action="store_true", help="For any keys in the reference but missing from " "the deepspeed checkpoint, copy them from the " "reference safetensors. Use when training has " "frozen modules (e.g. llm_adapter_lr=0) — those " "weights aren't saved in checkpoints but are " "needed for inference. With this flag, the output " "is a complete model = trained_modules ∪ frozen_from_ref.") return p.parse_args() def load_reference_keys(ref_path: Path) -> dict[str, torch.Size]: """Return {key: shape} from the reference safetensors.""" ref: dict[str, torch.Size] = {} with safe_open(str(ref_path), framework="pt", device="cpu") as f: for k in f.keys(): t = f.get_tensor(k) ref[k] = t.shape return ref def discover_layers(ckpt_dir: Path) -> list[Path]: """Return sorted list of layer_NN-model_states.pt paths.""" files = sorted(ckpt_dir.glob("layer_*-model_states.pt")) if not files: print(f"ERROR: no layer_*-model_states.pt under {ckpt_dir}", file=sys.stderr) sys.exit(2) return files def load_layer_state(path: Path) -> dict[str, torch.Tensor]: """Load one layer file. Returns its state_dict (empty dict if empty).""" sd = torch.load(str(path), map_location="cpu", weights_only=False) if not isinstance(sd, dict): print(f"WARN unexpected type in {path.name}: {type(sd).__name__}", file=sys.stderr) return {} # Filter to only tensors (skip metadata if any). return {k: v for k, v in sd.items() if isinstance(v, torch.Tensor)} def convert(ckpt_dir: Path) -> OrderedDict[str, torch.Tensor]: """ Walk all layer .pt files in order, remap keys to Anima's namespace. Strategy: - Bucket layers into [embedding, block_0, block_1, ..., final]. - Empty layers are skipped (passthrough). - The FIRST non-empty layer with x_embedder/t_embedder/etc keys becomes the embedding layer (prepend "net."). - The LAST non-empty layer with final_layer.* keys becomes the final layer (prepend "net."). - Everything in between with "block.*" keys is a transformer block, rename to "net.blocks.." where i counts up from 0 across them. """ layer_files = discover_layers(ckpt_dir) print(f"found {len(layer_files)} layer files in {ckpt_dir.name}") # First pass: load every layer, classify it. loaded: list[tuple[Path, dict[str, torch.Tensor]]] = [] for f in layer_files: sd = load_layer_state(f) loaded.append((f, sd)) # Classify embedding_layers: list[dict[str, torch.Tensor]] = [] block_layers: list[dict[str, torch.Tensor]] = [] final_layers: list[dict[str, torch.Tensor]] = [] empty_count = 0 for f, sd in loaded: if not sd: empty_count += 1 continue # Determine type by inspecting top-level keys any_block = any(k.startswith("block.") for k in sd.keys()) any_final = any(k.startswith("final_layer.") for k in sd.keys()) any_embed = any(k.startswith(p) for k in sd.keys() for p in ("x_embedder.", "t_embedder.", "t_embedding_norm.")) if any_block: block_layers.append(sd) elif any_final: final_layers.append(sd) elif any_embed: embedding_layers.append(sd) else: print(f"WARN unclassifiable layer {f.name}: top keys = " f"{list(sd.keys())[:3]}", file=sys.stderr) print(f" embedding layers: {len(embedding_layers)}") print(f" block layers: {len(block_layers)}") print(f" final layers: {len(final_layers)}") print(f" empty (passthrough) layers: {empty_count}") if len(embedding_layers) != 1: print(f"WARN expected 1 embedding layer, found {len(embedding_layers)}", file=sys.stderr) if len(final_layers) != 1: print(f"WARN expected 1 final layer, found {len(final_layers)}", file=sys.stderr) # Build the output state dict with Anima's "net." prefix. out: OrderedDict[str, torch.Tensor] = OrderedDict() # Embedding (prepend "net.") for sd in embedding_layers: for k, v in sd.items(): out[f"net.{k}"] = v # Blocks (rename "block." → "net.blocks..") for i, sd in enumerate(block_layers): prefix = f"net.blocks.{i}." for k, v in sd.items(): if k.startswith("block."): out[prefix + k[len("block."):]] = v else: print(f"WARN block {i} has unexpected key '{k}'", file=sys.stderr) out[prefix + k] = v # Final (prepend "net.") for sd in final_layers: for k, v in sd.items(): out[f"net.{k}"] = v return out def verify(out: dict[str, torch.Tensor], ref: dict[str, torch.Size]) -> bool: """Return True iff converted output matches reference (keys + shapes).""" ok = True out_keys = set(out.keys()) ref_keys = set(ref.keys()) missing_in_out = ref_keys - out_keys extra_in_out = out_keys - ref_keys if missing_in_out: print(f"FAIL {len(missing_in_out)} keys missing in converted output:", file=sys.stderr) for k in sorted(missing_in_out)[:10]: print(f" - {k}", file=sys.stderr) if len(missing_in_out) > 10: print(f" ... and {len(missing_in_out) - 10} more", file=sys.stderr) ok = False if extra_in_out: print(f"FAIL {len(extra_in_out)} keys in converted output that aren't in reference:", file=sys.stderr) for k in sorted(extra_in_out)[:10]: print(f" + {k}", file=sys.stderr) if len(extra_in_out) > 10: print(f" ... and {len(extra_in_out) - 10} more", file=sys.stderr) ok = False # Shape verification on overlapping keys shape_mismatches: list[tuple[str, tuple, tuple]] = [] for k in out_keys & ref_keys: out_shape = tuple(out[k].shape) ref_shape = tuple(ref[k]) if out_shape != ref_shape: shape_mismatches.append((k, out_shape, ref_shape)) if shape_mismatches: print(f"FAIL {len(shape_mismatches)} shape mismatches:", file=sys.stderr) for k, o, r in shape_mismatches[:10]: print(f" {k}: converted={o} vs reference={r}", file=sys.stderr) if len(shape_mismatches) > 10: print(f" ... and {len(shape_mismatches) - 10} more", file=sys.stderr) ok = False if ok: print(f"OK {len(out_keys)} keys match reference exactly") print(f"OK all tensor shapes match") return ok def main() -> int: args = parse_args() if not args.deepspeed_dir.is_dir(): print(f"--deepspeed-dir not a directory: {args.deepspeed_dir}", file=sys.stderr) return 2 if not args.reference_safetensors.is_file(): print(f"--reference-safetensors missing: {args.reference_safetensors}", file=sys.stderr) return 3 target_dtype = { "bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32, }[args.dtype] print(f"loading reference key/shape index from {args.reference_safetensors.name} ...") t0 = time.monotonic() ref = load_reference_keys(args.reference_safetensors) print(f" {len(ref)} keys in {time.monotonic() - t0:.1f}s") print() print(f"converting {args.deepspeed_dir.name} ...") t0 = time.monotonic() out = convert(args.deepspeed_dir) print(f" converted {len(out)} keys in {time.monotonic() - t0:.1f}s") # Merge in frozen modules from the reference (e.g. llm_adapter when # llm_adapter_lr=0 was set in training). DeepSpeed only saves trained # parameters, so frozen ones are missing from the checkpoint — but # they're identical to base Anima by construction (no gradient updates). if args.merge_frozen_from_ref: missing_in_out = set(ref.keys()) - set(out.keys()) if missing_in_out: print() print(f"merging {len(missing_in_out)} frozen keys from reference " f"(--merge-frozen-from-ref) ...") top_prefixes = set() for k in missing_in_out: top_prefixes.add(".".join(k.split(".")[:2])) print(f" top-level prefixes being merged: {sorted(top_prefixes)}") with safe_open(str(args.reference_safetensors), framework="pt", device="cpu") as f: for k in missing_in_out: out[k] = f.get_tensor(k) print(f" merged. output now has {len(out)} keys.") # Cast all tensors to target dtype if target_dtype != torch.float32: print(f"casting to {args.dtype} ...") out = OrderedDict((k, v.to(target_dtype)) for k, v in out.items()) # Verify if not args.skip_verify: print() print("verifying ...") if not verify(out, ref): print("verification FAILED — not writing output", file=sys.stderr) return 1 # Save args.output.parent.mkdir(parents=True, exist_ok=True) print() print(f"writing {args.output} ...") t0 = time.monotonic() # Make all tensors contiguous (safetensors requires it) out = OrderedDict((k, v.contiguous()) for k, v in out.items()) save_file(out, str(args.output)) sz = args.output.stat().st_size / (1024 ** 3) print(f" wrote {sz:.2f} GB in {time.monotonic() - t0:.1f}s") return 0 if __name__ == "__main__": sys.exit(main())