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