| |
| """Generate manifest file with MD5 checksums for model weights. |
| |
| Usage: |
| python -m tools.generate_manifest ckpts/Z-Image-Turbo |
| python -m tools.generate_manifest ckpts/Z-Image-Turbo --no-checksums # Only list files |
| """ |
|
|
| import argparse |
| import hashlib |
| from pathlib import Path |
| from typing import List |
|
|
|
|
| def compute_md5(file_path: Path, chunk_size: int = 8192) -> str: |
| """Compute MD5 hash of a file.""" |
| md5_hash = hashlib.md5() |
| with open(file_path, "rb") as f: |
| while chunk := f.read(chunk_size): |
| md5_hash.update(chunk) |
| return md5_hash.hexdigest() |
|
|
|
|
| def get_essential_files(model_dir: Path) -> List[Path]: |
| """Get list of essential model files.""" |
| essential_patterns = [ |
| "model_index.json", |
| "transformer/config.json", |
| "transformer/*.safetensors*", |
| "vae/config.json", |
| "vae/*.safetensors", |
| "text_encoder/config.json", |
| "text_encoder/*.safetensors*", |
| "tokenizer/tokenizer.json", |
| "tokenizer/tokenizer_config.json", |
| "scheduler/scheduler_config.json", |
| ] |
| |
| files = [] |
| for pattern in essential_patterns: |
| if "*" in pattern: |
| files.extend(model_dir.glob(pattern)) |
| else: |
| file_path = model_dir / pattern |
| if file_path.exists(): |
| files.append(file_path) |
| |
| return sorted(files) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Generate manifest file for model weights") |
| parser.add_argument("model_dir", type=str, help="Path to model directory") |
| parser.add_argument("--output", "-o", type=str, default=None, |
| help="Output manifest file path (default: auto-detect to config/manifests/)") |
| parser.add_argument("--no-checksums", action="store_true", |
| help="Only list files without computing checksums") |
| parser.add_argument("--verbose", "-v", action="store_true", |
| help="Print progress") |
| |
| args = parser.parse_args() |
| |
| model_dir = Path(args.model_dir) |
| if not model_dir.exists(): |
| print(f"Error: Model directory not found: {model_dir}") |
| return 1 |
| |
| |
| if args.output: |
| output_file = Path(args.output) |
| else: |
| |
| model_name = model_dir.name.lower() |
| script_dir = Path(__file__).parent |
| config_dir = script_dir.parent / "config" / "manifests" |
| config_dir.mkdir(parents=True, exist_ok=True) |
| output_file = config_dir / f"{model_name}.txt" |
| |
| |
| files = get_essential_files(model_dir) |
| |
| if not files: |
| print(f"Warning: No essential files found in {model_dir}") |
| return 1 |
| |
| print(f"Found {len(files)} essential files") |
| |
| |
| with open(output_file, "w", encoding="utf-8") as f: |
| f.write("# Z-Image Model Manifest\n") |
| if args.no_checksums: |
| f.write("# Format: <filepath>\n") |
| else: |
| f.write("# Format: <md5hash> <filepath>\n") |
| f.write("# Generated automatically - DO NOT edit manually\n\n") |
| |
| for file_path in files: |
| rel_path = file_path.relative_to(model_dir) |
| |
| if args.no_checksums: |
| f.write(f"{rel_path}\n") |
| if args.verbose: |
| print(f" {rel_path}") |
| else: |
| if args.verbose: |
| print(f"Computing MD5 for {rel_path}...", end=" ", flush=True) |
| |
| try: |
| md5_hash = compute_md5(file_path) |
| f.write(f"{md5_hash} {rel_path}\n") |
| if args.verbose: |
| print(f"✓ {md5_hash}") |
| except Exception as e: |
| print(f"✗ Error: {e}") |
| continue |
| |
| print(f"\n✓ Manifest saved to: {output_file}") |
| print(f" Total files: {len(files)}") |
| if not args.no_checksums: |
| print(f" With MD5 checksums for integrity verification") |
| |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| exit(main()) |
|
|
|
|
