""" PRISM DiT Benchmark & Evaluation Suite. Executes standard evaluation runs across the 5 core domains demonstrated in the PRISM showcase. """ import json import os import time from typing import Dict, Any import torch from model import PrismDiT from pipeline import PrismPipeline BENCHMARK_PROMPTS = [ { "id": "01_abstract_mandala", "category": "Geometric Abstraction & Prismatic Optics", "prompt": "kaleidoscopic radial mandala, abstract geometric human figures arranged in a circle, sharp faceted crystalline shapes, chromatic prism refraction, vivid spectral rainbow gradients, deep contrast, clean sharp edges", "negative_prompt": "blurry, muddy colors, desaturated, photorealistic, noise, low resolution", "height": 512, "width": 512, "steps": 18, "cfg": 6.5, "seed": 42, }, { "id": "02_waterfall_landscape", "category": "Cinematic Nature & Fluid Dynamics", "prompt": "long exposure photograph of a majestic waterfall cascading down volcanic basalt cliffs, turquoise glacial plunge pool, dark moss-covered river stones in foreground, misty atmosphere, soft cinematic daylight, 8k resolution, National Geographic style", "negative_prompt": "cartoon, oversaturated, harsh direct flash, CGI, digital painting, blown highlights", "height": 288, "width": 512, "steps": 18, "cfg": 6.0, "seed": 108, }, { "id": "03_indie_animation", "category": "Traditional 2D Storybook Animation", "prompt": "storyboard illustration, 2D indie animation style, a female animator with dark bob hair sitting at a drawing desk with drafting paper, loose expressive pencil outlines, subtle watercolor gouache wash, warm vintage paper texture, charming character design", "negative_prompt": "3D render, glossy, photorealistic, octane render, airbrushed, plastic skin", "height": 512, "width": 512, "steps": 18, "cfg": 5.5, "seed": 256, }, { "id": "04_vintage_coupe", "category": "Automotive & Chromatic Dispersion", "prompt": "vintage classic sedan automobile parked on dry plain, low angle 3/4 front view, chrome bumper reflections, vivid prismatic rainbow light refraction caustics sweeping across side fender, dreamlike atmospheric finish", "negative_prompt": "distorted wheels, asymmetrical car, low poly, oversaturated artifacts", "height": 512, "width": 512, "steps": 18, "cfg": 6.0, "seed": 512, }, { "id": "05_fine_art_lion", "category": "Traditional Dry Medium Wildlife Portraiture", "prompt": "fine art wildlife portrait of a roaring male lion walking forward, fierce snarling expression, bared teeth, detailed dense mane fur, drawn with chalk pastel and colored pencil on textured toned craft paper, traditional art medium", "negative_prompt": "digital vector, 3D model, smooth plastic, photograph, oversaturated, deformed eyes", "height": 512, "width": 384, "steps": 18, "cfg": 6.5, "seed": 1024, }, ] def run_benchmark( checkpoint_path: str = None, output_dir: str = "eval_outputs", device: str = "cpu", fast_check: bool = False, ): """Executes the 5 benchmark tests and saves summary logs.""" os.makedirs(output_dir, exist_ok=True) print(f"=== Starting PRISM DiT Benchmark Suite (Device: {device}) ===") # Initialize a lightweight DiT model for evaluation testing # If checkpoint_path is provided, real weights will be loaded dit = PrismDiT( input_size=64, patch_size=2, in_channels=16, hidden_size=384 if fast_check else 768, depth=4 if fast_check else 12, num_heads=6 if fast_check else 12, context_dim=1024, ) if checkpoint_path and os.path.exists(checkpoint_path): print(f"Loading weights from {checkpoint_path}...") if checkpoint_path.endswith(".safetensors"): from safetensors.torch import load_file state_dict = load_file(checkpoint_path) else: state_dict = torch.load(checkpoint_path, map_location="cpu") dit.load_state_dict(state_dict, strict=False) else: print("Using initialized model weights for architecture verification...") pipe = PrismPipeline(dit=dit, device=device) results = [] for test in BENCHMARK_PROMPTS: test_id = test["id"] steps = 3 if fast_check else test["steps"] h = test["height"] w = test["width"] print(f"\nRunning [{test_id}] - {test['category']} ({w}x{h}, {steps} steps)...") start_time = time.time() images = pipe( prompt=test["prompt"], negative_prompt=test["negative_prompt"], height=h, width=w, num_inference_steps=steps, guidance_scale=test["cfg"], seed=test["seed"], ) elapsed = time.time() - start_time out_image_path = os.path.join(output_dir, f"{test_id}.png") images[0].save(out_image_path) print(f"Saved output to {out_image_path} ({elapsed:.2f}s)") results.append({ "test_id": test_id, "category": test["category"], "prompt": test["prompt"], "seed": test["seed"], "resolution": f"{w}x{h}", "steps": steps, "elapsed_seconds": round(elapsed, 2), "output_file": out_image_path, }) summary_file = os.path.join(output_dir, "benchmark_summary.json") with open(summary_file, "w") as f: json.dump(results, f, indent=2) print(f"\n=== Benchmark Complete. Summary written to {summary_file} ===") return results if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="PRISM DiT Benchmark Runner") parser.add_argument("--checkpoint", type=str, default=None, help="Path to .safetensors or .pt checkpoint") parser.add_argument("--output-dir", type=str, default="eval_outputs", help="Output directory for generated images") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") parser.add_argument("--fast-check", action="store_true", help="Run shallow verification pass") args = parser.parse_args() run_benchmark( checkpoint_path=args.checkpoint, output_dir=args.output_dir, device=args.device, fast_check=args.fast_check, )