File size: 4,719 Bytes
7c325a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """
PRISM Checkpoint Exporter & Serializer.
Packages the Custom DiT model into standard PRISM-1.safetensors with embedded metadata.
"""
import argparse
import os
import time
from typing import Dict, Any
import torch
from safetensors.torch import save_file, load_file
from model import PrismDiT
def export_prism_safetensors(
output_path: str = "PRISM-1.safetensors",
precision: str = "fp16",
hidden_size: int = 1152,
depth: int = 28,
num_heads: int = 16,
context_dim: int = 1024,
source_weights_path: str = None,
) -> str:
"""
Initializes or loads PRISM DiT weights and serializes to .safetensors with complete metadata.
"""
print(f"=== Initializing PRISM-1 Checkpoint Generator ===")
print(f"Configuration:")
print(f" - Hidden Size: {hidden_size}")
print(f" - Depth (Layers): {depth}")
print(f" - Attention Heads: {num_heads}")
print(f" - Context Dimension: {context_dim}")
print(f" - Precision: {precision}")
print(f" - Target File: {output_path}")
# 1. Instantiate DiT Model
model = PrismDiT(
input_size=64,
patch_size=2,
in_channels=16,
hidden_size=hidden_size,
depth=depth,
num_heads=num_heads,
context_dim=context_dim,
)
# 2. Optionally load existing source weights
if source_weights_path and os.path.exists(source_weights_path):
print(f"Loading source weights from: {source_weights_path}")
if source_weights_path.endswith(".safetensors"):
state_dict = load_file(source_weights_path)
else:
state_dict = torch.load(source_weights_path, map_location="cpu")
if "state_dict" in state_dict:
state_dict = state_dict["state_dict"]
elif "model" in state_dict:
state_dict = state_dict["model"]
model.load_state_dict(state_dict, strict=False)
print("Successfully merged source weights into PRISM DiT.")
else:
print("No external source weights provided; generating clean base initialized checkpoint...")
# 3. Cast precision
dtype_map = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}
target_dtype = dtype_map.get(precision.lower(), torch.float16)
model = model.to(dtype=target_dtype)
# 4. Extract contiguous state dict
state_dict = {k: v.contiguous() for k, v in model.state_dict().items()}
total_params = sum(p.numel() for p in state_dict.values())
print(f"Total Parameters: {total_params / 1e6:.2f}M")
# 5. Metadata for Hugging Face / Civitai / ComfyUI loader inspection
metadata: Dict[str, str] = {
"model_name": "PRISM-1",
"architecture": "PrismDiT (Diffusion Transformer)",
"version": "1.0.0",
"format": "safetensors",
"precision": precision,
"in_channels": "16",
"patch_size": "2",
"hidden_size": str(hidden_size),
"depth": str(depth),
"num_heads": str(num_heads),
"context_dim": str(context_dim),
"diffusion_type": "Flow-Matching / Rectified Flow",
"recommended_sampler": "Euler",
"recommended_steps": "12-25",
"recommended_cfg": "5.0-6.5",
"generator": "PRISM Checkpoint Exporter v1.0",
}
# 6. Save safetensors
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
t0 = time.time()
save_file(state_dict, output_path, metadata=metadata)
elapsed = time.time() - t0
file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"\n[SUCCESS] Checkpoint saved successfully!")
print(f" File: {output_path}")
print(f" Size: {file_size_mb:.2f} MB")
print(f" Time taken: {elapsed:.2f}s")
return output_path
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Export PRISM-1.safetensors Checkpoint")
parser.add_argument("--output", type=str, default="PRISM-1.safetensors")
parser.add_argument("--precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--hidden-size", type=int, default=1152)
parser.add_argument("--depth", type=int, default=28)
parser.add_argument("--heads", type=int, default=16)
parser.add_argument("--context-dim", type=int, default=1024)
parser.add_argument("--source-weights", type=str, default=None)
args = parser.parse_args()
export_prism_safetensors(
output_path=args.output,
precision=args.precision,
hidden_size=args.hidden_size,
depth=args.depth,
num_heads=args.heads,
context_dim=args.context_dim,
source_weights_path=args.source_weights,
)
|