diff --git a/README.md b/README.md index 89e336a..ccf4845 100644 --- a/README.md +++ b/README.md @@ -139,6 +139,43 @@ Either model can be omitted (`None`) for single-model sessions. Properties: The standalone `vnasset generate` and `vnasset edit` CLI commands are thin wrappers around a one-shot session — same API, backwards compatible. +### Background Removal + +Remove backgrounds from character sprites (output is RGBA PNG): + +```bash +# Single file +vnasset remove-bg --input character_base.png --output character_transparent.png + +# Batch (reuses model across files) +vnasset remove-bg --input base.png --input happy.png --input sad.png --output-dir transparent/ +``` + +Or via the session API: + +```python +with VnAssetsSession() as vna: + vna.remove_background("base.png", output="base_transparent.png") + + # Batch + vna.remove_backgrounds( + ["base.png", "happy.png", "sad.png"], + output_dir="transparent/", + ) +``` + +| Option | Default | Description | +|--------|---------|-------------| +| `--input` | (required, repeatable) | Input image path(s) | +| `--output` | (auto) | Output path (single mode) | +| `--output-dir` | (none) | Output directory (batch mode) | +| `--model` | `isnet-anime` | Model: `isnet-anime`, `u2net`, `u2netp`, `u2net_human_seg`, `isnet-general-use`, `sam` | + +The default `isnet-anime` model is trained on anime images — ideal for the +novaAnimeXL art style. Inference runs on CPU via onnxruntime (~0.25s per +1024×1024 image on Strix Halo). The model is ~176 MB, downloaded on first use +to `~/.u2net/`. + ## Architecture ``` @@ -350,7 +387,8 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding. | Compel prompt weighting + BREAK | ✅ Working | | Lightning LoRA fuse-at-load | ✅ Working | | Flash attention (experimental) | ✅ Working | -| Metadata JSON output | ✅ Working | +| `vnasset remove-bg` | ✅ Working | +| Session background removal | ✅ Working | | `vnasset pipeline` (batch YAML config) | 🚧 Planned | | `vnasset serve` (daemon/HTTP API) | 🚧 Planned | | `torch.compile` on UNet | 🚧 Planned | @@ -371,5 +409,3 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding. availability or a ComfyUI installation. - **`vnasset serve`** — lightweight daemon with Unix socket or HTTP API for integrating VNAsset into external tools. -- **Background removal** — green-screen trim and transparency pass for - post-processing sprites. diff --git a/pyproject.toml b/pyproject.toml index b79a66a..f2a9293 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,8 @@ dependencies = [ "pyyaml", "click", "compel", + "rembg", + "onnxruntime", ] [tool.setuptools.packages.find] diff --git a/vnassets/background.py b/vnassets/background.py new file mode 100644 index 0000000..0f8d58e --- /dev/null +++ b/vnassets/background.py @@ -0,0 +1,99 @@ +"""Background removal for character sprites using rembg + isnet-anime. + +Uses onnxruntime on CPU — the model is ~176 MB and runs in ~0.25s per +1024×1024 image on Strix Halo. No GPU dependency keeps things simple. +""" +import time +from pathlib import Path +from typing import Sequence + +from PIL import Image +from rembg import new_session, remove + + +def _load_session(model: str = "isnet-anime"): + """Load a rembg session for the given model. Cached in the caller.""" + return new_session(model) + + +def remove_background( + input_path: str, + output_path: str, + model: str = "isnet-anime", + session=None, +) -> None: + """Remove background from a single image. Output is RGBA PNG. + + Args: + input_path: Path to the input RGB image. + output_path: Where to save the RGBA PNG. + model: rembg model name. Default "isnet-anime" is trained on anime + images. Other options: "u2net", "u2netp", "u2net_human_seg", + "isnet-general-use", "sam". + session: A rembg session (from ``new_session()``). If None, one + is created and immediately discarded. For batch use, create a + session once and pass it to avoid reloading the model. + + Raises: + FileNotFoundError: If input_path doesn't exist. + """ + input_file = Path(input_path) + if not input_file.exists(): + raise FileNotFoundError(f"Input image not found: {input_path}") + + output_file = Path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + + _session = session + close_session = session is None + if _session is None: + _session = _load_session(model) + + img = Image.open(input_file).convert("RGB") + + t0 = time.perf_counter() + result = remove(img, session=_session) + t_elapsed = round(time.perf_counter() - t0, 3) + + result.save(output_file) + print(f"Saved {output_file} (bg removed, {t_elapsed}s)") + + if close_session: + del _session + + +def remove_backgrounds( + input_paths: Sequence[str], + output_dir: str, + model: str = "isnet-anime", +) -> None: + """Remove backgrounds from multiple images, reusing one model session. + + Output files are named ``{stem}_nobg.png`` in ``output_dir``. + + Args: + input_paths: List of input image paths. + output_dir: Directory for output RGBA PNGs. + model: rembg model name. Default "isnet-anime". + """ + out_dir = Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + session = _load_session(model) + try: + total = 0.0 + for input_path in input_paths: + stem = Path(input_path).stem + output_path = out_dir / f"{stem}_nobg.png" + + img = Image.open(input_path).convert("RGB") + t0 = time.perf_counter() + result = remove(img, session=session) + t_elapsed = time.perf_counter() - t0 + total += t_elapsed + result.save(output_path) + print(f"Saved {output_path} ({t_elapsed:.3f}s)") + + print(f"Done: {len(input_paths)} images, {total:.2f}s total") + finally: + del session diff --git a/vnassets/cli.py b/vnassets/cli.py index 0a3886d..ff630f3 100644 --- a/vnassets/cli.py +++ b/vnassets/cli.py @@ -1,6 +1,8 @@ """VNAsset — Visual Novel Asset Pipeline CLI.""" 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 @@ -79,3 +81,42 @@ def edit_cmd(model, input_path, prompt, steps, cfg, seed, output, lora_path): output_path=output, lora_path=lora_path, ) + + +REMOVE_BG_MODELS = ["isnet-anime", "u2net", "u2netp", "u2net_human_seg", "isnet-general-use", "sam"] + + +@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 + from pathlib import Path + 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." + ) diff --git a/vnassets/session.py b/vnassets/session.py index 4380628..b297f4b 100644 --- a/vnassets/session.py +++ b/vnassets/session.py @@ -23,6 +23,8 @@ from diffusers.models.transformers import QwenImageTransformer2DModel from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor from .attention import patch_qwen_transformer, patch_unet_attention +from .background import remove_background as _remove_bg +from .background import remove_backgrounds as _remove_bgs from .prompt import build_compel, encode_prompts TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct" @@ -69,6 +71,7 @@ class VnAssetsSession: self._pipe_sdxl: StableDiffusionXLPipeline | None = None self._compel = None self._pipe_qwen: QwenImageEditPlusPipeline | None = None + self._rembg_session = None t0 = time.perf_counter() if sdxl_checkpoint: @@ -339,6 +342,55 @@ class VnAssetsSession: meta_path.write_text(json.dumps(meta, indent=2)) print(f"Saved {meta_path}") + # ── Background Removal ───────────────────────────────────────────── + + def remove_background( + self, + input_path: str, + output_path: str, + model: str = "isnet-anime", + ) -> None: + """Remove background from an image. Output is RGBA PNG. + + Wraps rembg with ``isnet-anime`` (trained on anime images) by + default. The rembg session is created once and reused across + calls, so batch use is efficient. + + Args: + input_path: Path to the input RGB image. + output_path: Where to save the RGBA PNG. + model: rembg model name. Default ``isnet-anime``. Also + available: ``u2net``, ``u2netp``, ``u2net_human_seg``, + ``isnet-general-use``, ``sam``. + + Raises: + FileNotFoundError: If input_path doesn't exist. + """ + if self._rembg_session is None: + from rembg import new_session + self._rembg_session = new_session(model) + _remove_bg(input_path, output_path, model=model, session=self._rembg_session) + + def remove_backgrounds( + self, + input_paths: list[str], + output_dir: str, + model: str = "isnet-anime", + ) -> None: + """Remove backgrounds from multiple images in one batch. + + Output files are named ``{stem}_nobg.png``. + + Args: + input_paths: List of input image paths. + output_dir: Directory for output RGBA PNGs. + model: rembg model name. Default ``isnet-anime``. + """ + if self._rembg_session is None: + from rembg import new_session + self._rembg_session = new_session(model) + _remove_bgs(input_paths, output_dir, model=model) + # ── Lifecycle ─────────────────────────────────────────────────────── def close(self) -> None: @@ -350,6 +402,9 @@ class VnAssetsSession: del self._pipe_qwen self._pipe_qwen = None self._compel = None + if self._rembg_session: + del self._rembg_session + self._rembg_session = None if self.device == "cuda": torch.cuda.empty_cache()