| """ |
| 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}") |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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...") |
|
|
| |
| 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) |
|
|
| |
| 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") |
|
|
| |
| 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", |
| } |
|
|
| |
| 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, |
| ) |
|
|
