src/qwen_image_edit_lite/cli.py

from __future__ import annotations

import argparse
from pathlib import Path

from qwen_image_edit_lite.backend import (
    DEFAULT_CONFIG_PATH,
    AppConfig,
    DEFAULT_DOWNLOAD_CACHE_DIR,
    DEFAULT_LORAS_DIR,
    InferenceRequest,
    ensure_valid_resolution,
    flatten_extra_args,
    load_config,
    normalize_lora_scales,
    resolve_path_from_base,
    run_inference,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Minimal Qwen Image Edit 2511 inference wrapper built on copied ComfyUI runtime code."
    )
    parser.add_argument("--input", nargs="+", required=True, help="One or more input images. image1 is required.")
    parser.add_argument("--prompt", required=True, help="Edit instruction.")
    parser.add_argument("--negative-prompt", default=None, help="Optional negative prompt.")
    parser.add_argument("--output-dir", default="outputs", help="Directory for generated images.")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG_PATH), help="Optional config.ini path.")

    parser.add_argument("--models-dir", default=None, help="Override base model directory.")
    parser.add_argument("--diffusion-models-dir", default=None, help="Override diffusion model directory.")
    parser.add_argument("--vae-dir", default=None, help="Override VAE directory.")
    parser.add_argument("--text-encoder-dir", default=None, help="Override text encoder directory.")
    parser.add_argument("--loras-dir", default=None, help="Override LoRA directory.")

    parser.add_argument("--lora", nargs="*", default=[], help="Zero or more Qwen Image Edit LoRA files.")
    parser.add_argument(
        "--lora-scale",
        nargs="*",
        type=float,
        default=None,
        help="LoRA multipliers. If omitted, all LoRA files use 1.0.",
    )

    parser.add_argument("--width", type=int, default=1024, help="Output width. Must be divisible by 16.")
    parser.add_argument("--height", type=int, default=1024, help="Output height. Must be divisible by 16.")
    parser.add_argument("--steps", type=int, default=30, help="Number of denoising steps.")
    parser.add_argument("--cfg", type=float, default=4.0, help="Classifier-free guidance scale.")
    parser.add_argument("--seed", type=int, default=None, help="Random seed.")
    parser.add_argument("--flow-shift", type=float, default=None, help="Optional AuraFlow shift override.")
    parser.add_argument("--mask", default=None, help="Optional inpainting mask path.")

    parser.add_argument("--device", default=None, help="Torch device, for example cuda, cuda:0, or cpu.")
    parser.add_argument("--attn-mode", choices=("torch", "sdpa", "flash", "xformers"), default=None)
    parser.add_argument("--disable-xformers-qwen-patch", action="store_true", help="Use stock ComfyUI xformers without the Lite Qwen compact patch.")
    parser.add_argument("--text-encoder-cpu", action="store_true", help="Run the text encoder on CPU.")
    parser.add_argument("--extra-arg", action="append", default=[], help="Reserved compatibility option.")
    return parser


def build_runtime_config(args: argparse.Namespace) -> AppConfig:
    base_config = load_config(args.config)
    project_root = base_config.project_root

    if args.models_dir:
        models_dir = resolve_path_from_base(args.models_dir, project_root, "models_dir")
        diffusion_models_dir = resolve_path_from_base("diffusion_models", models_dir, "diffusion_models")
        vae_dir = resolve_path_from_base("vae", models_dir, "vae")
        text_encoders_dir = resolve_path_from_base("text_encoders", models_dir, "text_encoders")
    else:
        diffusion_models_dir = base_config.diffusion_models_dir
        vae_dir = base_config.vae_dir
        text_encoders_dir = base_config.text_encoders_dir

    if args.diffusion_models_dir:
        diffusion_models_dir = resolve_path_from_base(args.diffusion_models_dir, project_root, "diffusion_models")
    if args.vae_dir:
        vae_dir = resolve_path_from_base(args.vae_dir, project_root, "vae")
    if args.text_encoder_dir:
        text_encoders_dir = resolve_path_from_base(args.text_encoder_dir, project_root, "text_encoders")

    loras_dir = base_config.loras_dir
    if args.loras_dir:
        loras_dir = resolve_path_from_base(args.loras_dir, project_root, "loras")

    return AppConfig(
        project_root=project_root,
        diffusion_models_ref=str(diffusion_models_dir),
        vae_ref=str(vae_dir),
        text_encoders_ref=str(text_encoders_dir),
        loras_dir=loras_dir,
        download_cache_dir=(project_root / DEFAULT_DOWNLOAD_CACHE_DIR).resolve(),
        hf_token=base_config.hf_token,
        text_encoder_cpu=args.text_encoder_cpu if args.text_encoder_cpu else base_config.text_encoder_cpu,
        device=args.device or base_config.device,
        attn_mode=args.attn_mode or base_config.attn_mode,
        xformers_qwen_patch=False if args.disable_xformers_qwen_patch else base_config.xformers_qwen_patch,
        default_steps=base_config.default_steps,
        default_cfg=base_config.default_cfg,
        default_flow_shift=base_config.default_flow_shift,
        presets=base_config.presets,
    )


def resolve_lora_paths(lora_args: list[str], loras_dir: Path) -> list[Path]:
    resolved = []
    for raw in lora_args:
        path = Path(raw).expanduser()
        if path.is_absolute():
            resolved.append(path.resolve())
            continue

        candidate = (loras_dir / path).resolve()
        if candidate.exists():
            resolved.append(candidate)
            continue

        resolved.append((Path.cwd() / path).resolve())
    return resolved


def main() -> None:
    args = build_parser().parse_args()
    ensure_valid_resolution(args.width, args.height)

    config = build_runtime_config(args)
    input_paths = [Path(raw).expanduser().resolve() for raw in args.input]
    lora_paths = resolve_lora_paths(args.lora, config.loras_dir)

    request = InferenceRequest(
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        control_images=input_paths,
        lora_paths=lora_paths,
        lora_scales=normalize_lora_scales(lora_paths, args.lora_scale),
        width=args.width,
        height=args.height,
        steps=args.steps,
        cfg=args.cfg,
        seed=args.seed,
        flow_shift=args.flow_shift,
        mask=Path(args.mask).expanduser().resolve() if args.mask else None,
        output_dir=Path(args.output_dir).expanduser().resolve(),
        extra_args=flatten_extra_args(args.extra_arg),
    )
    run_inference(config, request)


if __name__ == "__main__":
    main()