Add ability to run small scripts from the ui and added a merge lora script

ui_scripts/merge_loras.py

@@ -0,0 +1,134 @@
+"""Merge a list of LoRAs into a single checkpoint."""
+
+import argparse
+import json
+import os
+import sys
+
+import torch
+from safetensors.torch import load_file, save_file
+from safetensors import safe_open
+
+
+DTYPE_MAP = {
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "float16": torch.float16,
+    "fp16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+}
+
+
+def log(message: str) -> None:
+    print(message, flush=True)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Merge a list of LoRAs into a single checkpoint."
+    )
+    parser.add_argument(
+        "--loras",
+        required=True,
+        help='JSON list of {"path": "...", "strength": 1.0} entries.',
+    )
+    parser.add_argument("--output", required=True, help="Output .safetensors path.")
+    parser.add_argument(
+        "--save_dtype",
+        default="bfloat16",
+        choices=list(DTYPE_MAP.keys()),
+        help="Dtype of the saved tensors (merging is always done in float32).",
+    )
+    parser.add_argument(
+        "--device",
+        default="cpu",
+        help="Device to merge on (cpu, cuda, cuda:1, mps). Output is always saved from CPU.",
+    )
+    args = parser.parse_args()
+
+    try:
+        loras = json.loads(args.loras)
+    except json.JSONDecodeError as e:
+        print(f"Failed to parse --loras JSON: {e}", file=sys.stderr, flush=True)
+        return 2
+
+    if not isinstance(loras, list) or len(loras) == 0:
+        print("--loras must be a non-empty JSON list.", file=sys.stderr, flush=True)
+        return 2
+
+    device = torch.device(args.device)
+    save_dtype = DTYPE_MAP[args.save_dtype]
+
+    log(f"Merging {len(loras)} LoRA(s) on {device}, saving as {args.save_dtype}.")
+
+    merged: dict[str, torch.Tensor] = {}
+
+    metadata = {}
+
+    for i, entry in enumerate(loras):
+        if not isinstance(entry, dict) or "path" not in entry:
+            print(
+                f"LoRA entry {i} must be an object with a 'path' field.",
+                file=sys.stderr,
+                flush=True,
+            )
+            return 2
+
+        path = entry["path"]
+        strength = float(entry.get("strength", 1.0))
+
+        if not os.path.isfile(path):
+            print(f"LoRA file not found: {path}", file=sys.stderr, flush=True)
+            return 2
+
+        log(f"[{i + 1}/{len(loras)}] Loading {path} (strength={strength})")
+        state_dict = load_file(path, device=str(device))
+
+        for key, tensor in state_dict.items():
+            scaled = tensor.to(torch.float32) * strength
+            if key in merged:
+                merged[key].add_(scaled)
+            else:
+                merged[key] = scaled
+        del state_dict
+
+        if i == 0:
+            # For the first LoRA, also copy over all non-tensor metadata (e.g. base model info)
+            with safe_open(path, framework="pt") as f:
+                metadata_to_keep = [
+                    "version",
+                    "format",
+                    "ss_base_model_version",
+                    "software",
+                ]
+                orig_metadata = f.metadata()
+                for meta_key in metadata_to_keep:
+                    if meta_key in orig_metadata:
+                        metadata[meta_key] = orig_metadata[meta_key]
+
+    log(f"Casting to {args.save_dtype} and moving to CPU")
+    final = {k: v.to(save_dtype).cpu().contiguous() for k, v in merged.items()}
+    merged.clear()
+
+    log(f"Saving merged checkpoint to {args.output}")
+    save_file(final, args.output, metadata=metadata)
+
+    print(
+        json.dumps(
+            {
+                "ok": True,
+                "output": args.output,
+                "num_loras": len(loras),
+                "num_keys": len(final),
+                "save_dtype": args.save_dtype,
+                "device": str(device),
+            }
+        ),
+        flush=True,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

ui_scripts/test_script.py

@@ -0,0 +1,46 @@
+"""Example UI script invoked via POST /api/scripts.
+
+The API streams stdout/stderr back to the caller. To return structured data,
+print a single JSON object on the final line of stdout — the route parses it
+into the `result` field of the response.
+"""
+
+import argparse
+import json
+import sys
+import time
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Example ui_scripts entry point.")
+    parser.add_argument("--message", default="hello", help="Message to echo back.")
+    parser.add_argument(
+        "--count", type=int, default=3, help="Number of log lines to emit."
+    )
+    parser.add_argument(
+        "--delay", type=float, default=0.0, help="Seconds to sleep between log lines."
+    )
+    parser.add_argument(
+        "--fail", action="store_true", help="Exit non-zero to demo failure handling."
+    )
+    args = parser.parse_args()
+
+    for i in range(args.count):
+        print(f"[{i + 1}/{args.count}] {args.message}", flush=True)
+        if args.delay > 0:
+            time.sleep(args.delay)
+
+    if args.fail:
+        print("intentional failure", file=sys.stderr, flush=True)
+        return 1
+
+    # Final line: JSON payload the API surfaces as `result`.
+    print(
+        json.dumps({"message": args.message, "count": args.count, "ok": True}),
+        flush=True,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())