vnassets/background: add bg removal via rembg + isnet-anime
rembg with ONNX Runtime on CPU was chosen over BiRefNet: the latter requires deform_conv2d which crashes on ROCm bfloat16 and runs at 40s in float32. rembg delivers 0.23s per 1024x1024 image, no GPU deps, and isnet-anime is trained specifically on anime images — exactly the target domain. CLI and session API both support single-image and batch (plural-first) modes, reusing one model session across files.
This commit is contained in:
42
README.md
42
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.
|
||||
|
||||
@@ -16,6 +16,8 @@ dependencies = [
|
||||
"pyyaml",
|
||||
"click",
|
||||
"compel",
|
||||
"rembg",
|
||||
"onnxruntime",
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
|
||||
99
vnassets/background.py
Normal file
99
vnassets/background.py
Normal file
@@ -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
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user