| |
| """ |
| Repository model inspector. |
| |
| This script is designed to work in `--config-only` mode without importing |
| PyTorch/Diffusers/Transformers. It reads JSON configs from a local Diffusers |
| repository layout and prints a summary. |
| |
| With `--params`, it can also compute parameter counts by scanning |
| `*.safetensors` headers (without loading tensor data into RAM). |
| """ |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
| from typing import Any, Dict, Iterable, Optional |
|
|
|
|
| def load_json(path: Path) -> Dict[str, Any]: |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def human_params(value: Optional[int]) -> str: |
| if value is None: |
| return "n/a" |
| if value >= 1_000_000_000: |
| return f"{value/1e9:.2f}B" |
| return f"{value/1e6:.2f}M" |
|
|
|
|
| def read_model_index(model_dir: Path) -> Dict[str, Any]: |
| idx_path = model_dir / "model_index.json" |
| if not idx_path.exists(): |
| return {} |
| return load_json(idx_path) |
|
|
|
|
| def describe_model_index(model_index: Dict[str, Any]) -> None: |
| if not model_index: |
| return |
| print("Pipeline pieces (model_index.json):") |
| for key, val in model_index.items(): |
| if key.startswith("_"): |
| continue |
| print(f" {key:14s} -> {val}") |
| print() |
|
|
|
|
| def detect_pipeline_kind(model_index: Dict[str, Any]) -> str: |
| cls = str(model_index.get("_class_name", "")).lower() |
| if "zimage" in cls or ("transformer" in model_index and "unet" not in model_index): |
| return "zimage" |
| if "stable" in cls or "unet" in model_index: |
| return "sdxl_like" |
| return "unknown" |
|
|
|
|
| def iter_safetensors_files(directory: Path) -> Iterable[Path]: |
| if not directory.exists(): |
| return [] |
| return sorted(p for p in directory.iterdir() if p.is_file() and p.suffix == ".safetensors") |
|
|
|
|
| def count_params_from_safetensors(files: Iterable[Path]) -> int: |
| from safetensors import safe_open |
|
|
| total = 0 |
| for file in files: |
| with safe_open(str(file), framework="np") as f: |
| for key in f.keys(): |
| shape = f.get_slice(key).get_shape() |
| total += math.prod(shape) |
| return int(total) |
|
|
|
|
| def zimage_config_only_summary(model_dir: Path, include_params: bool) -> Dict[str, Any]: |
| model_index = read_model_index(model_dir) |
|
|
| te_cfg_path = model_dir / "text_encoder" / "config.json" |
| transformer_cfg_path = model_dir / "transformer" / "config.json" |
| vae_cfg_path = model_dir / "vae" / "config.json" |
| scheduler_cfg_path = model_dir / "scheduler" / "scheduler_config.json" |
|
|
| te_cfg = load_json(te_cfg_path) if te_cfg_path.exists() else {} |
| transformer_cfg = load_json(transformer_cfg_path) if transformer_cfg_path.exists() else {} |
| vae_cfg = load_json(vae_cfg_path) if vae_cfg_path.exists() else {} |
| scheduler_cfg = load_json(scheduler_cfg_path) if scheduler_cfg_path.exists() else {} |
|
|
| text_encoder_params = None |
| transformer_params = None |
| vae_params = None |
| if include_params: |
| text_encoder_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "text_encoder")) |
| transformer_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "transformer")) |
| vae_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "vae")) |
|
|
| print("[Text encoder]") |
| if te_cfg: |
| arch = te_cfg.get("architectures", []) |
| arch_name = arch[0] if isinstance(arch, list) and arch else "n/a" |
| print(f" architecture={arch_name}") |
| print( |
| " " |
| f"layers={te_cfg.get('num_hidden_layers', 'n/a')}, " |
| f"hidden={te_cfg.get('hidden_size', 'n/a')}, " |
| f"heads={te_cfg.get('num_attention_heads', 'n/a')}, " |
| f"intermediate={te_cfg.get('intermediate_size', 'n/a')}" |
| ) |
| print(f" vocab={te_cfg.get('vocab_size', 'n/a')}, max_positions={te_cfg.get('max_position_embeddings', 'n/a')}") |
| else: |
| print(" [warn] missing text_encoder/config.json") |
| print(f" params={human_params(text_encoder_params)}") |
| print() |
|
|
| print("[Transformer]") |
| if transformer_cfg: |
| print(f" class={transformer_cfg.get('_class_name', 'n/a')}") |
| print( |
| " " |
| f"dim={transformer_cfg.get('dim', 'n/a')}, " |
| f"layers={transformer_cfg.get('n_layers', 'n/a')}, " |
| f"heads={transformer_cfg.get('n_heads', 'n/a')}" |
| ) |
| print(f" in_channels={transformer_cfg.get('in_channels', 'n/a')}, cap_feat_dim={transformer_cfg.get('cap_feat_dim', 'n/a')}") |
| print(f" patch_size={transformer_cfg.get('all_patch_size', 'n/a')}, f_patch_size={transformer_cfg.get('all_f_patch_size', 'n/a')}") |
| else: |
| print(" [warn] missing transformer/config.json") |
| print(f" params={human_params(transformer_params)}") |
| print() |
|
|
| print("[VAE]") |
| if vae_cfg: |
| print(f" class={vae_cfg.get('_class_name', 'n/a')}") |
| print( |
| " " |
| f"sample_size={vae_cfg.get('sample_size', 'n/a')}, " |
| f"in_channels={vae_cfg.get('in_channels', 'n/a')}, " |
| f"latent_channels={vae_cfg.get('latent_channels', 'n/a')}, " |
| f"out_channels={vae_cfg.get('out_channels', 'n/a')}" |
| ) |
| print(f" block_out_channels={vae_cfg.get('block_out_channels', 'n/a')}, scaling_factor={vae_cfg.get('scaling_factor', 'n/a')}") |
| else: |
| print(" [warn] missing vae/config.json") |
| print(f" params={human_params(vae_params)}") |
| print() |
|
|
| print("[Scheduler]") |
| if scheduler_cfg: |
| print( |
| " " |
| f"class={scheduler_cfg.get('_class_name', 'n/a')}, " |
| f"timesteps={scheduler_cfg.get('num_train_timesteps', 'n/a')}, " |
| f"shift={scheduler_cfg.get('shift', 'n/a')}" |
| ) |
| else: |
| print(" [warn] missing scheduler/scheduler_config.json") |
| print() |
|
|
| return { |
| "kind": "zimage", |
| "pipeline": model_index, |
| "text_encoder": {"config": te_cfg, "params": text_encoder_params}, |
| "transformer": {"config": transformer_cfg, "params": transformer_params}, |
| "vae": {"config": vae_cfg, "params": vae_params}, |
| "scheduler": {"config": scheduler_cfg}, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Inspect a local Diffusers-style repository layout.") |
| parser.add_argument("--model-dir", type=Path, default=Path(".."), help="Path to the diffusers pipeline directory.") |
| parser.add_argument("--device", default="cpu", help="Unused (kept for CLI compatibility).") |
| parser.add_argument("--fp16", action="store_true", help="Unused (kept for CLI compatibility).") |
| parser.add_argument("--config-only", action="store_true", help="Read JSON configs and print a summary.") |
| parser.add_argument("--params", action="store_true", help="Count parameters from *.safetensors headers (no tensor loading).") |
| parser.add_argument("--json-out", type=Path, default=None, help="Write a JSON summary to this path.") |
| args = parser.parse_args() |
|
|
| model_index = read_model_index(args.model_dir) |
| if not model_index: |
| raise SystemExit(f"model_index.json not found under {args.model_dir}") |
|
|
| describe_model_index(model_index) |
| kind = detect_pipeline_kind(model_index) |
|
|
| if not args.config_only: |
| raise SystemExit("Only --config-only mode is supported by this inspector.") |
|
|
| if kind != "zimage": |
| raise SystemExit(f"Unsupported pipeline kind: {kind} (expected ZImagePipeline-style layout)") |
|
|
| summary = zimage_config_only_summary(args.model_dir, include_params=args.params) |
|
|
| if args.json_out is not None: |
| args.json_out.parent.mkdir(parents=True, exist_ok=True) |
| args.json_out.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") |
| print(f"[info] wrote JSON summary to {args.json_out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
