| |
| """Run one Mage-Flow Edit-Turbo XPO3 image edit.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from pathlib import Path |
| import sys |
| import time |
|
|
| from PIL import Image |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| for path in (ROOT / "runtime", ROOT / "vendor"): |
| sys.path.insert(0, str(path)) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("reference", type=Path) |
| parser.add_argument("instruction") |
| parser.add_argument("--output", type=Path, default=Path("edited.png")) |
| parser.add_argument("--seed", type=int, default=1) |
| parser.add_argument("--max-size", type=int, default=1024) |
| parser.add_argument("--disable-fused-gelu-up", action="store_true") |
| parser.add_argument("--disable-fp4-bridge", action="store_true") |
| parser.add_argument("--disable-accelerated-attention", action="store_true") |
| parser.add_argument("--disable-direct-hnd", action="store_true") |
| parser.add_argument("--report", type=Path) |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") |
|
|
| import torch |
| from portable_turbo_runtime import ( |
| close_pipeline_optimization_runtimes, |
| generation_optimization_context, |
| load_pipeline_from_files, |
| transformer_config, |
| ) |
|
|
| if not torch.cuda.is_available(): |
| raise RuntimeError("XPO3 requires an NVIDIA CUDA GPU") |
| torch.cuda.set_device(0) |
| properties = torch.cuda.get_device_properties(0) |
| if (properties.major, properties.minor) != (12, 0): |
| raise RuntimeError( |
| "this XPO3 build requires an NVIDIA Blackwell SM120 GPU; " |
| f"found compute capability {properties.major}.{properties.minor}" |
| ) |
|
|
| diffusion_model = ( |
| ROOT / "Mage-Flow-Edit-Turbo-XPO3-NVFP4.safetensors" |
| ) |
| text_encoder = ROOT / "qwen3vl_4b_fp8_scaled.safetensors" |
| vae = ROOT / "Mage-Flow-VAE.safetensors" |
| config = transformer_config(diffusion_model) |
| profile = config["quantization_config"]["xpo3_runtime_profile"] |
| bridge_blocks = ",".join( |
| sorted( |
| profile["fp4_bridge_scales"], |
| key=lambda value: int(value), |
| ) |
| ) |
| attention_steps = ",".join( |
| str(value) for value in profile["attention_steps"] |
| ) |
| attention_blocks = ",".join( |
| str(value) for value in profile["attention_blocks"] |
| ) |
|
|
| torch.manual_seed(args.seed) |
| torch.cuda.manual_seed_all(args.seed) |
| torch.cuda.reset_peak_memory_stats(0) |
| started = time.perf_counter() |
| pipe, load_report = load_pipeline_from_files( |
| diffusion_model=diffusion_model, |
| text_encoder=text_encoder, |
| vae=vae, |
| support_root=ROOT / "resources" / "turbo", |
| fused_gelu_library=ROOT / "runtime/libmage_nvfp4_gelu_up.so", |
| bridge_up_library=ROOT / "runtime/libmage_nvfp4_bridge_up.so", |
| bridge_down_library=( |
| ROOT / "runtime/libmage_nvfp4_prequantized_down.so" |
| ), |
| torch=torch, |
| ) |
| for module in ( |
| pipe.model.txt_enc, |
| pipe.model.transformer, |
| pipe.model.vae, |
| ): |
| module.to("cuda:0") |
| torch.cuda.synchronize() |
| load_seconds = time.perf_counter() - started |
|
|
| reference = Image.open(args.reference).convert("RGB") |
| generation_started = time.perf_counter() |
| with torch.inference_mode(): |
| with generation_optimization_context( |
| pipe=pipe, |
| torch=torch, |
| enable_fused_gelu_up=not args.disable_fused_gelu_up, |
| enable_fp4_bridge=not args.disable_fp4_bridge, |
| bridge_blocks=bridge_blocks, |
| enable_attention_accel=( |
| not args.disable_accelerated_attention |
| ), |
| enable_direct_hnd=not args.disable_direct_hnd, |
| attention_steps=attention_steps, |
| attention_blocks=attention_blocks, |
| steps=int(profile["steps"]), |
| static_shift=float(profile["static_shift"]), |
| cfg=float(profile["cfg"]), |
| ) as feature_manifest: |
| images = pipe.edit( |
| [args.instruction], |
| [[reference]], |
| seeds=[args.seed], |
| steps=int(profile["steps"]), |
| cfg=float(profile["cfg"]), |
| max_size=args.max_size, |
| static_shift=float(profile["static_shift"]), |
| prompt_template="mage-flow-edit", |
| vl_cond_long_edge=384, |
| ) |
| torch.cuda.synchronize() |
| generation_seconds = time.perf_counter() - generation_started |
| if not feature_manifest["restoration"]["all_restored"]: |
| raise RuntimeError("XPO3 feature state did not restore after editing") |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| images[0].save(args.output) |
| report = { |
| "schema_version": "mage-flow-edit-turbo-xpo3-cli-v1", |
| "reference": str(args.reference.resolve()), |
| "instruction": args.instruction, |
| "output": str(args.output.resolve()), |
| "seed": args.seed, |
| "max_size": args.max_size, |
| "gpu": { |
| "name": properties.name, |
| "compute_capability": ( |
| f"{properties.major}.{properties.minor}" |
| ), |
| }, |
| "load_and_placement_seconds": load_seconds, |
| "generation_seconds": generation_seconds, |
| "peak_allocated_bytes": int(torch.cuda.max_memory_allocated(0)), |
| "load": load_report, |
| "active_feature_manifest": feature_manifest, |
| } |
| if args.report is not None: |
| args.report.parent.mkdir(parents=True, exist_ok=True) |
| args.report.write_text( |
| json.dumps(report, indent=2, sort_keys=True) + "\n", |
| encoding="utf-8", |
| ) |
| print(json.dumps(report, indent=2, sort_keys=True)) |
| cleanup = close_pipeline_optimization_runtimes(pipe) |
| if not cleanup["all_closed_without_error"]: |
| raise RuntimeError(f"XPO3 native cleanup failed: {cleanup}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
