studio/cli.py

"""studio CLI — verbs mirror cursor ops.

Subcommands:
    open    <image>                                 validate; print cursor state
    compile <kf0> <kf1> -o <out.npz>                two-keyframe timing-layer compile
    export  <frames.npz> -o <out.gif|.mp4|.webm>    render FrameStack to disk
"""
from __future__ import annotations
import argparse
import sys

import json

from pixel_cursor import open_cursor, ease_between, load_image
from . import Studio
from .io import save_framestack_npz, load_framestack_npz, export_video


def cmd_open(args: argparse.Namespace) -> int:
    img = load_image(args.image)
    cur = open_cursor(img)
    print(f"opened: {img}")
    for k, v in cur.tell().items():
        print(f"  {k}: {v}")
    return 0


def cmd_compile(args: argparse.Namespace) -> int:
    kf0 = load_image(args.kf0)
    kf1 = load_image(args.kf1)
    cur_a = open_cursor(kf0).at_frame(0)
    cur_b = open_cursor(kf1).at_frame(args.frames - 1)
    seg = ease_between(0, args.frames, curve=args.curve, stride=args.stride)
    stack = Studio().compile_24fps_from_keyframes([cur_a, cur_b], [seg])
    save_framestack_npz(stack, args.out)
    print(f"compiled {stack.num_frames} frames @ {stack.fps}fps -> {args.out}")
    print(f"  curve={args.curve} stride={args.stride} (on-{['ones','twos','threes'][min(args.stride-1, 2)]})")
    return 0


def cmd_export(args: argparse.Namespace) -> int:
    stack = load_framestack_npz(args.frames)
    out = export_video(stack, args.out)
    print(f"exported {stack.num_frames} frames @ {stack.fps}fps -> {out}")
    return 0


def cmd_sheet(args: argparse.Namespace) -> int:
    from .spritesheet import write_sprite_sheet, sheet_metadata

    stack = load_framestack_npz(args.frames)
    out = write_sprite_sheet(stack, args.out, cols=args.cols)
    meta = sheet_metadata(stack, cols=args.cols)
    print(f"sheet -> {out}")
    print(json.dumps(meta, indent=2))
    if args.manifest:
        manifest_path = args.manifest
        with open(manifest_path, "w") as f:
            json.dump(meta, f, indent=2)
        print(f"manifest -> {manifest_path}")
    return 0


def cmd_motion(args: argparse.Namespace) -> int:
    from .backends.animatediff import AnimateDiffAdapter, MOTION_LORA_MAP
    from .backends.hf_space import HFSpaceAdapter

    if args.list_presets:
        print("MotionLoRA presets (shared by `animatediff` and `hf_space` backends):")
        for name, lora_id in MOTION_LORA_MAP.items():
            print(f"  {name:24s} -> {lora_id}")
        return 0

    if args.preset is None or args.image is None:
        print(
            "error: --preset and --image are required (or use --list-presets)",
            file=sys.stderr,
        )
        return 2

    if args.backend == "hf_space":
        if not args.space:
            print("error: --space <owner>/<name> is required for --backend hf_space", file=sys.stderr)
            return 2
        adapter = HFSpaceAdapter(space_id=args.space, hf_token=args.hf_token)
    else:
        adapter = AnimateDiffAdapter()
    adapter.register()

    img = load_image(args.image)
    cur = open_cursor(img).bind_motion_exemplar([args.preset], backend=args.backend)

    motion_spec = {
        "num_frames": args.num_frames,
        "num_inference_steps": args.steps,
        "guidance_scale": args.guidance,
        "prompt": args.prompt,
        "negative_prompt": args.negative_prompt,
        "seed": args.seed,
    }

    if args.dry_run:
        diag = adapter.dry_run(cur, motion_spec)
        print(json.dumps(diag, indent=2, default=str))
        return 0

    if args.out is None:
        print("error: -o/--out is required unless --dry-run is set", file=sys.stderr)
        return 2

    stack = cur.write_motion(motion_spec, backend=args.backend)
    save_framestack_npz(stack, args.out)
    print(f"generated {stack.num_frames} frames @ {stack.fps}fps -> {args.out}")
    print(f"  backend={args.backend}  preset={args.preset}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="studio",
        description="Pixel-cursor animation studio (v0: model-free timing compile)",
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    p_open = sub.add_parser("open", help="Load an image; print cursor state")
    p_open.add_argument("image")
    p_open.set_defaults(func=cmd_open)

    p_compile = sub.add_parser(
        "compile",
        help="Compile a two-keyframe segment into a FrameStack via the cursor timing layer",
    )
    p_compile.add_argument("kf0", help="first keyframe image")
    p_compile.add_argument("kf1", help="second keyframe image")
    p_compile.add_argument("-o", "--out", required=True, help="output .npz path")
    p_compile.add_argument("--frames", type=int, default=24, help="total frames (default 24)")
    p_compile.add_argument(
        "--curve",
        default="ease",
        choices=("linear", "ease", "ease-in", "ease-out", "spring"),
    )
    p_compile.add_argument(
        "--stride",
        type=int,
        default=2,
        help="on-twos=2 (default), on-ones=1, on-threes=3",
    )
    p_compile.add_argument("--fps", type=int, default=24, help="target fps (default 24)")
    p_compile.set_defaults(func=cmd_compile)

    p_export = sub.add_parser(
        "export", help="Render a FrameStack .npz to a video file"
    )
    p_export.add_argument("frames", help="input .npz from `studio compile`")
    p_export.add_argument(
        "-o", "--out", required=True, help="output .gif | .mp4 | .webm | .mkv | .avi"
    )
    p_export.set_defaults(func=cmd_export)

    p_sheet = sub.add_parser(
        "sheet",
        help="Pack a FrameStack .npz into a sprite-sheet PNG (Babylon SpriteManager friendly)",
    )
    p_sheet.add_argument("frames", help="input .npz from `studio compile` or `studio motion`")
    p_sheet.add_argument("-o", "--out", required=True, help="output .png path")
    p_sheet.add_argument("--cols", type=int, default=4, help="frames per row in the grid (default 4)")
    p_sheet.add_argument("--manifest", help="optional .json manifest path (cell size, fps, loop ms)")
    p_sheet.set_defaults(func=cmd_sheet)

    p_motion = sub.add_parser(
        "motion",
        help="Generate motion on a single image via AnimateDiff MotionLoRA preset",
    )
    p_motion.add_argument(
        "--list-presets", action="store_true", help="List all 8 verified MotionLoRA presets and exit"
    )
    p_motion.add_argument("--preset", help="preset name (see --list-presets)")
    p_motion.add_argument("--image", help="source image path")
    p_motion.add_argument("-o", "--out", help="output .npz path (required unless --dry-run)")
    p_motion.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would happen + resource estimate; no model load, no output written",
    )
    p_motion.add_argument("--num-frames", type=int, default=16)
    p_motion.add_argument("--steps", type=int, default=25, help="num_inference_steps")
    p_motion.add_argument("--guidance", type=float, default=7.5, help="guidance_scale")
    p_motion.add_argument("--prompt", default="high quality, detailed")
    p_motion.add_argument("--negative-prompt", default="bad quality, blurry")
    p_motion.add_argument("--seed", type=int, default=42)
    p_motion.add_argument(
        "--backend",
        default="animatediff",
        choices=("animatediff", "hf_space"),
        help="local diffusers run (animatediff) or remote ZeroGPU Space (hf_space)",
    )
    p_motion.add_argument(
        "--space", default=None,
        help="HF Space id (e.g. owner/venture-studio); required when --backend hf_space",
    )
    p_motion.add_argument(
        "--hf-token", default=None,
        help="HF token (overrides HF_TOKEN env var and ~/.cache/huggingface/token)",
    )
    p_motion.set_defaults(func=cmd_motion)

    return p


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())