Add edit.py

edit.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""Run one Mage-Flow Edit 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-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" / "edit",
+        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"]),
+            required_cfg=float(profile["cfg"]),
+            expected_steps=int(profile["steps"]),
+        ) as feature_manifest:
+            images = pipe.edit(
+                [args.instruction],
+                [[reference]],
+                neg_prompts=[" "],
+                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,
+                batch_cfg=False,
+            )
+    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-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())