Add microscope config-only inspector

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+.DS_Store
+.venv/
+microscope/.venv/
+microscope/__pycache__/

microscope/README.md

@@ -0,0 +1,42 @@
+# Microscope tooling
+
+Small utilities to inspect this repository’s Diffusers checkpoint without leaving the repo.
+
+## Setup
+
+From the repo root:
+
+```bash
+cd microscope
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -U pip
+pip install -r requirements.txt
+```
+
+## Usage
+
+Inspect the local model (default: one directory up from this folder) without importing PyTorch/Diffusers:
+
+```bash
+python inspect_model.py --model-dir .. --config-only
+```
+
+Include parameter counts by scanning `*.safetensors` headers (still does not load weights into RAM):
+
+```bash
+python inspect_model.py --model-dir .. --config-only --params
+```
+
+Write a machine-readable summary:
+
+```bash
+python inspect_model.py --model-dir .. --config-only --params --json-out model_summary.json
+```
+
+Flags:
+
+- `--model-dir`: path to the Diffusers pipeline (default: `..`).
+- `--config-only`: read JSON configs and print a summary (recommended).
+- `--params`: count parameters from `*.safetensors` metadata (no tensor loading).
+- `--json-out`: write a JSON summary to this path.

microscope/inspect_model.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+"""
+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()

microscope/model_summary.json

@@ -0,0 +1,144 @@
+{
+  "kind": "zimage",
+  "pipeline": {
+    "_class_name": "ZImagePipeline",
+    "_diffusers_version": "0.36.0.dev0",
+    "scheduler": [
+      "diffusers",
+      "FlowMatchEulerDiscreteScheduler"
+    ],
+    "text_encoder": [
+      "transformers",
+      "Qwen3Model"
+    ],
+    "tokenizer": [
+      "transformers",
+      "Qwen2Tokenizer"
+    ],
+    "transformer": [
+      "diffusers",
+      "ZImageTransformer2DModel"
+    ],
+    "vae": [
+      "diffusers",
+      "AutoencoderKL"
+    ]
+  },
+  "text_encoder": {
+    "config": {
+      "architectures": [
+        "Qwen3ForCausalLM"
+      ],
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "bos_token_id": 151643,
+      "eos_token_id": 151645,
+      "head_dim": 128,
+      "hidden_act": "silu",
+      "hidden_size": 2560,
+      "initializer_range": 0.02,
+      "intermediate_size": 9728,
+      "max_position_embeddings": 40960,
+      "max_window_layers": 36,
+      "model_type": "qwen3",
+      "num_attention_heads": 32,
+      "num_hidden_layers": 36,
+      "num_key_value_heads": 8,
+      "rms_norm_eps": 1e-06,
+      "rope_scaling": null,
+      "rope_theta": 1000000,
+      "sliding_window": null,
+      "tie_word_embeddings": true,
+      "torch_dtype": "bfloat16",
+      "transformers_version": "4.51.0",
+      "use_cache": true,
+      "use_sliding_window": false,
+      "vocab_size": 151936
+    },
+    "params": 4022468096
+  },
+  "transformer": {
+    "config": {
+      "_class_name": "ZImageTransformer2DModel",
+      "_diffusers_version": "0.36.0.dev0",
+      "all_f_patch_size": [
+        1
+      ],
+      "all_patch_size": [
+        2
+      ],
+      "axes_dims": [
+        32,
+        48,
+        48
+      ],
+      "axes_lens": [
+        1536,
+        512,
+        512
+      ],
+      "cap_feat_dim": 2560,
+      "dim": 3840,
+      "in_channels": 16,
+      "n_heads": 30,
+      "n_kv_heads": 30,
+      "n_layers": 30,
+      "n_refiner_layers": 2,
+      "norm_eps": 1e-05,
+      "qk_norm": true,
+      "rope_theta": 256.0,
+      "t_scale": 1000.0
+    },
+    "params": 6154908736
+  },
+  "vae": {
+    "config": {
+      "_class_name": "AutoencoderKL",
+      "_diffusers_version": "0.36.0.dev0",
+      "_name_or_path": "flux-dev",
+      "act_fn": "silu",
+      "block_out_channels": [
+        128,
+        256,
+        512,
+        512
+      ],
+      "down_block_types": [
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D"
+      ],
+      "force_upcast": true,
+      "in_channels": 3,
+      "latent_channels": 16,
+      "latents_mean": null,
+      "latents_std": null,
+      "layers_per_block": 2,
+      "mid_block_add_attention": true,
+      "norm_num_groups": 32,
+      "out_channels": 3,
+      "sample_size": 1024,
+      "scaling_factor": 0.3611,
+      "shift_factor": 0.1159,
+      "up_block_types": [
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D"
+      ],
+      "use_post_quant_conv": false,
+      "use_quant_conv": false
+    },
+    "params": 83819683
+  },
+  "scheduler": {
+    "config": {
+      "_class_name": "FlowMatchEulerDiscreteScheduler",
+      "_diffusers_version": "0.36.0.dev0",
+      "num_train_timesteps": 1000,
+      "use_dynamic_shifting": false,
+      "shift": 3.0
+    }
+  }
+}

microscope/requirements.txt

@@ -0,0 +1 @@
+safetensors>=0.4.2