vnassets/vnassets/cli.py

"""VNAsset — Visual Novel Asset Pipeline CLI."""
from pathlib import Path

import click

from .background import remove_background as _remove_bg
from .background import remove_backgrounds as _remove_bgs
from .edit import edit
from .generate import generate
from .upscale import upscale as _upscale_fn
from .upscale import upscales as _upscales_fn


def _parse_seed(ctx, param, value):
    if value is None or value == "":
        return None
    if value.lower() == "random":
        return None
    try:
        return int(value)
    except ValueError:
        raise click.BadParameter("seed must be an integer or 'random'")


@click.group()
def main():
    """VNAsset — fast CLI pipeline for visual novel image assets."""


@main.command()
@click.option("--checkpoint", required=True, help="Path to SDXL checkpoint (.safetensors)")
@click.option("--prompt", required=True, help="Positive prompt")
@click.option("--negative-prompt", default="", help="Negative prompt")
@click.option("--width", default=1024, type=int)
@click.option("--height", default=1024, type=int)
@click.option("--steps", default=20, type=int)
@click.option("--cfg", default=4.5, type=float)
@click.option(
    "--seed", default="random", callback=_parse_seed,
    help="RNG seed (integer or 'random')",
)
@click.option("--output", default="output.png", help="Output image path")
@click.option(
    "--raw", is_flag=True,
    help="Disable ComfyUI-style prompt weighting (use plain diffusers encoding)",
)
def generate_cmd(checkpoint, prompt, negative_prompt, width, height, steps, cfg, seed, output, raw):
    """Generate an image from an SDXL checkpoint."""
    generate(
        checkpoint_path=checkpoint,
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        steps=steps,
        cfg=cfg,
        seed=seed,
        output_path=output,
        raw=raw,
    )


@main.command()
@click.option("--model", required=True, help="Path to Qwen Image Edit diffusion model (.safetensors)")
@click.option("--input", "input_path", required=True, help="Input image to edit")
@click.option("--prompt", required=True, help="Edit instruction")
@click.option("--steps", default=20, type=int, help="Inference steps")
@click.option("--cfg", default=4.0, type=float, help="CFG scale")
@click.option(
    "--seed", default="random", callback=_parse_seed,
    help="RNG seed (integer or 'random')",
)
@click.option("--output", default="output.png", help="Output image path")
@click.option("--lora", "lora_path", default=None,
              help="Path to LoRA .safetensors (e.g. Lightning 4-step LoRA)")
def edit_cmd(model, input_path, prompt, steps, cfg, seed, output, lora_path):
    """Edit an image using Qwen Image Edit."""
    edit(
        model_path=model,
        input_path=input_path,
        prompt=prompt,
        steps=steps,
        cfg=cfg,
        seed=seed,
        output_path=output,
        lora_path=lora_path,
    )


REMOVE_BG_MODELS = ["isnet-anime", "u2net", "u2netp", "u2net_human_seg", "isnet-general-use", "sam"]

UPSALE_SCALES = [2, 3, 4]


@main.command("upscale")
@click.option("--input", "input_paths", multiple=True, required=True,
              help="Input image path (repeat for batch).")
@click.option("--output", "output_path", default=None,
              help="Output path (single-input mode).")
@click.option("--output-dir", default=None,
              help="Output directory (batch mode; files named {stem}_x{scale}.png).")
@click.option("--scale", default=2, type=click.Choice(["2", "3", "4"]),
              help="Upscale factor (default: 2).")
def upscale_cmd(input_paths, output_path, output_dir, scale):
    """Upscale images using Real-ESRGAN (anime model).

    Single file mode:

        vnasset upscale --input char.png --output char_2x.png --scale 2

    Batch mode (reuses model across files):

        vnasset upscale --input base.png --input happy.png --output-dir upscaled/ --scale 2
    """
    scale = int(scale)
    if len(input_paths) == 1 and output_path:
        _upscale_fn(input_paths[0], output_path, scale=scale)
    elif output_dir:
        _upscales_fn(list(input_paths), output_dir, scale=scale)
    elif len(input_paths) == 1:
        # Single input, no output specified: auto-name
        stem = Path(input_paths[0]).stem
        out = Path(input_paths[0]).parent / f"{stem}_x{scale}.png"
        _upscale_fn(input_paths[0], str(out), scale=scale)
    else:
        raise click.UsageError(
            "For multiple inputs, use --output-dir. For single input, use --output."
        )


@main.command("remove-bg")
@click.option("--input", "input_paths", multiple=True, required=True,
              help="Input image path (repeat for batch).")
@click.option("--output", "output_path", default=None,
              help="Output path (single-input mode).")
@click.option("--output-dir", default=None,
              help="Output directory (batch mode; files named {stem}_nobg.png).")
@click.option("--model", default="isnet-anime", type=click.Choice(REMOVE_BG_MODELS),
              help="Background removal model (default: isnet-anime).")
def remove_bg_cmd(input_paths, output_path, output_dir, model):
    """Remove background from images. Output is RGBA PNG.

    Single file mode:

        vnasset remove-bg --input char.png --output char_nobg.png

    Batch mode (reuses model across files):

        vnasset remove-bg --input base.png --input happy.png --output-dir transparent/
    """
    if len(input_paths) == 1 and output_path:
        _remove_bg(input_paths[0], output_path, model=model)
    elif output_dir:
        _remove_bgs(list(input_paths), output_dir, model=model)
    elif len(input_paths) == 1:
        # Single input, no output specified: auto-name
        stem = Path(input_paths[0]).stem
        out = Path(input_paths[0]).parent / f"{stem}_nobg.png"
        _remove_bg(input_paths[0], str(out), model=model)
    else:
        raise click.UsageError(
            "For multiple inputs, use --output-dir. For single input, use --output."
        )