jmp-games/vnassets
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add Real-ESRGAN upscaling (2x/3x/4x) with anime model

Browse Source 
Adds `vnasset upscale` CLI command, `VnAssetsSession.upscale()` / `.upscales()`
session methods, and a standalone `vnassets.upscale` module following the
existing remove-bg pattern.

Uses Real-ESRGAN RRDBNet with RealESRGAN_x4plus_anime_6B (~17 MB,
auto-downloaded). A 256x256 warmup tile at load time eliminates ~80s of
first-run CUDA JIT compilation on ROCm. Steady-state: ~1.8s per 1024->2048
upscale on Strix Halo. The upsampler is lazy-loaded on first call and coexists
with SDXL/Qwen in the same session.

Works around a basicsr/torchvision API incompatibility (rgb_to_grayscale moved
in torchvision 0.20+) with a 3-line module shim in upscale.py.
This commit is contained in:
Michele Rossi
2026-07-08 12:36:02 +02:00
parent 8a11325b2f
commit 23135e62cb
4 changed files with 365 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
README.md
View File

@@ -176,6 +176,46 @@ 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/`.
### Upscaling (Real-ESRGAN)
Upscale images 2×, 3×, or 4× using Real-ESRGAN with the anime-optimized
`RealESRGAN_x4plus_anime_6B` model:
```bash
# Single file
vnasset upscale --input character_base.png --output character_2x.png --scale 2
# Batch (reuses model across files)
vnasset upscale --input base.png --input happy.png --input sad.png --output-dir upscaled/ --scale 2
```
Or via the session API:
```python
with VnAssetsSession() as vna:
vna.upscale("base.png", output="base_2x.png", scale=2)
# Batch
vna.upscales(
["base.png", "happy.png", "sad.png"],
output_dir="upscaled/",
scale=2,
)
```
| Option | Default | Description |
|--------|---------|-------------|
| `--input` | (required, repeatable) | Input image path(s) |
| `--output` | (auto) | Output path (single mode) |
| `--output-dir` | (none) | Output directory (batch mode; `{stem}_x{scale}.png`) |
| `--scale` | `2` | Upscale factor: `2`, `3`, or `4` |
The model is ~17 MB (auto-downloaded from GitHub on first use). The upsampler
runs in FP32 on GPU and is lazy-loaded on first call within a session — it
coexists with SDXL and Qwen without memory pressure. A 256×256 warmup tile is
run at load time to compile CUDA kernels, so first-user upscale doesn't incur
a compilation penalty.
## Architecture
```
@@ -206,6 +246,19 @@ to `~/.u2net/`.
│ │ happy.png │ │
│ │ sad.png │ │
│ │ angry.png │ │
│ └────────┬─────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────┐ │
│ │ Upscale │ │
│ │ (Real-ESRGAN) │ │
│ │ ~1.8s each │ │
│ └────────┬─────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────┐ │
│ │ Remove BG │ │
│ │ (isnet-anime) │ │
│ └──────────────────┘ │
└─────────────────────────────────────────────────────────────┘
```
@@ -219,6 +272,10 @@ SDXL and Qwen use **separate VAEs** with different latent spaces. The SDXL
checkpoint bundles its own VAE; Qwen uses `Qwen/Qwen-Image` VAE from
HuggingFace. Both coexist in the same session without conflict.
The upsampler (Real-ESRGAN RRDBNet, ~17 MB, FP32) runs on GPU via PyTorch
and is lazy-loaded on first use. It uses tiled processing (512×512 tiles)
to keep memory modest even at 4× output (4096×4096).
## Data Flow
### Generate Phase
@@ -242,6 +299,18 @@ input image ──► Qwen VL 7B TE ──► visual tok.┘
The text encoder handles both text conditioning and visual token encoding from
the input image.
### Upscale Phase
```
input image (RGB) ──► tile split (512×512) ──┐
├──► RRDBNet (FP32, 4×) ──► tiles ──► blend ──► upscaled image
tile overlap padding ─────────────────────────┘
```
If the requested output scale is less than the model's native scale (e.g. 2×
from a 4× model), the 4× result is Lanczos-downsampled to the target size.
Tiled processing keeps GPU memory constant regardless of output resolution.
### Turbo vs Normal Mode
| Parameter | Normal | Turbo (Lightning LoRA) |
@@ -291,13 +360,32 @@ processor, not SDPA).
| Session load | ~28 | SDXL + Qwen transformer + VAE + TE + LoRA fuse |
| Generate | ~31 | SDXL 20-step, 1024×1024 |
| Edit (turbo) | ~87 / ~54 | Matmul / flash attention |
| **Total (2 ops)** | **~146 / ~118** | One session load amortized |
| Upscale (2×) | ~1.8 | Real-ESRGAN, 1024→2048 |
| **Total (3 ops)** | **~147 / ~119** | One session load amortized |
With a session, each additional edit saves one model-load round trip (~28s).
5 edits save 112s. Flash attention adds a further 1.6× multiplier on inference.
See [`docs/stats.md`](docs/stats.md) for detailed per-run breakdowns.
### Upscale (Real-ESRGAN, anime model)
| Scale | Input | Tiles | Load (s) | Inference (s) |
|-------|-------|-------|----------|---------------|
| 2× | 1024×1024 | 4 | ~0.2¹ | ~1.8 |
| 4× | 1024×1024 | 4 | ~0.2¹ | ~1.8 |
¹ Warmup tile only; model download is ~17 MB (first use, cached thereafter).
Upscaling is tiled (512×512 input tiles) to keep GPU memory modest — the
RRDBNet forward pass processes each tile independently, and the results are
blended at tile boundaries. The 2× and 4× paths have near-identical latency
because the model is natively 4×; 2× output is achieved by upscaling to 4×
then downsampling.
A 256×256 warmup tile is run at model load time to compile CUDA kernels.
Without it, the first real upscale incurs ~80s of JIT compilation.
## Technical Notes
### bfloat16 required on RDNA 3.5
@@ -389,6 +477,8 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding.
| Flash attention (experimental) | ✅ Working |
| `vnasset remove-bg` | ✅ Working |
| Session background removal | ✅ Working |
| `vnasset upscale` | ✅ Working |
| Session upscaling | ✅ Working |
| `vnasset pipeline` (batch YAML config) | 🚧 Planned |
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
| `torch.compile` on UNet | 🚧 Planned |

43
vnassets/cli.py
View File

@@ -1,10 +1,14 @@
"""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):
@@ -85,6 +89,44 @@ def edit_cmd(model, input_path, prompt, steps, cfg, seed, output, 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,
@@ -112,7 +154,6 @@ def remove_bg_cmd(input_paths, output_path, output_dir, model):
_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)

54
vnassets/session.py
View File

@@ -26,6 +26,9 @@ 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
from .upscale import _build_upsampler
from .upscale import upscale as _upscale_fn
from .upscale import upscales as _upscales_fn
TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
VAE_ID = "Qwen/Qwen-Image"
@@ -72,6 +75,7 @@ class VnAssetsSession:
self._compel = None
self._pipe_qwen: QwenImageEditPlusPipeline | None = None
self._rembg_session = None
self._upsampler = None
t0 = time.perf_counter()
if sdxl_checkpoint:
@@ -391,6 +395,53 @@ class VnAssetsSession:
self._rembg_session = new_session(model)
_remove_bgs(input_paths, output_dir, model=model)
# ── Upscaling ───────────────────────────────────────────────────
def upscale(
self,
input_path: str,
output_path: str,
scale: int = 2,
) -> None:
"""Upscale an image using Real-ESRGAN (anime model).
The upsampler is lazy-loaded on first call and reused across
calls. Uses ``RealESRGAN_x4plus_anime_6B`` (17 MB,
auto-downloaded).
Args:
input_path: Path to the input RGB image.
output_path: Where to save the upscaled PNG.
scale: Output scale factor (2, 3, or 4).
Raises:
FileNotFoundError: If input_path doesn't exist.
ValueError: If scale is not 2, 3, or 4.
"""
if self._upsampler is None:
self._upsampler = _build_upsampler(self.device)
# Also use this upsampler for subsequent calls
_upscale_fn(input_path, output_path, scale=scale, upsampler=self._upsampler)
def upscales(
self,
input_paths: list[str],
output_dir: str,
scale: int = 2,
) -> None:
"""Upscale multiple images, reusing one model session.
Output files are named ``{stem}_x{scale}.png``.
Args:
input_paths: List of input image paths.
output_dir: Directory for output PNGs.
scale: Output scale factor (2, 3, or 4).
"""
if self._upsampler is None:
self._upsampler = _build_upsampler(self.device)
_upscales_fn(input_paths, output_dir, scale=scale, upsampler=self._upsampler)
# ── Lifecycle ───────────────────────────────────────────────────────
def close(self) -> None:
@@ -405,6 +456,9 @@ class VnAssetsSession:
if self._rembg_session:
del self._rembg_session
self._rembg_session = None
if self._upsampler:
del self._upsampler
self._upsampler = None
if self.device == "cuda":
torch.cuda.empty_cache()

178
vnassets/upscale.py Normal file
View File

@@ -0,0 +1,178 @@
"""Image upscaling using Real-ESRGAN with the anime-optimized RRDBNet model.
Uses ``RealESRGAN_x4plus_anime_6B`` (17 MB, auto-downloaded on first use).
Output scale is configurable: 2x, 3x, or 4x. The model is natively 4x;
2x and 3x are achieved by upscaling to 4x then downsampling.
Follows the same pattern as ``background.py``: standalone functions that
can be called directly, with an optional persistent session for batch use.
"""
import sys
import time
import types
from pathlib import Path
from typing import Sequence
import numpy as np
import torch
from PIL import Image
# ── basicsr / torchvision compatibility shim ───────────────────────
# basicsr.data.degradations imports rgb_to_grayscale from the old
# torchvision path (functional_tensor), removed in torchvision 0.20+.
# The function lives in torchvision.transforms.functional now.
# We only need RRDBNet inference — this shim pacifies the import.
from torchvision.transforms import functional as _F_tv
_ft_mod = types.ModuleType("torchvision.transforms.functional_tensor")
_ft_mod.rgb_to_grayscale = _F_tv.rgb_to_grayscale
sys.modules["torchvision.transforms.functional_tensor"] = _ft_mod
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan.utils import RealESRGANer
# ── Constants ──────────────────────────────────────────────────────
ANIME_MODEL_URL = (
"https://github.com/xinntao/Real-ESRGAN/releases/download/"
"v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
)
ANIME_MODEL_NET_SCALE = 4
_WARMUP_SIZE = 256 # px — small tile to compile CUDA kernels at load time
def _build_upsampler(device: str = "cuda") -> RealESRGANer:
"""Create a RealESRGANer with the anime model.
Runs a tiny warmup forward pass to compile CUDA kernels, so the
first user-facing upscale doesn't pay an ~80s compilation penalty.
"""
model = RRDBNet(
num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4
)
upsampler = RealESRGANer(
scale=ANIME_MODEL_NET_SCALE,
model_path=ANIME_MODEL_URL,
model=model,
tile=512,
tile_pad=10,
pre_pad=0,
half=False,
device=device,
)
# Warmup: a 256×256 dummy tile compiles all CUDA kernels needed for
# the RRDBNet forward pass. Without this, the first real upscale of
# a 1024×1024 image incurs ~80s of JIT compilation.
dummy = np.zeros((_WARMUP_SIZE, _WARMUP_SIZE, 3), dtype=np.uint8)
upsampler.enhance(dummy, outscale=ANIME_MODEL_NET_SCALE)
if device == "cuda":
torch.cuda.synchronize()
return upsampler
def upscale(
input_path: str,
output_path: str,
scale: int = 2,
upsampler: RealESRGANer | None = None,
) -> None:
"""Upscale a single image. Output is an RGB PNG.
Args:
input_path: Path to the input image.
output_path: Where to save the upscaled PNG.
scale: Output scale factor (2, 3, or 4).
upsampler: A pre-built RealESRGANer. If None, one is created and
immediately discarded. For batch use, create one once and
pass it to avoid reloading the model.
Raises:
FileNotFoundError: If input_path doesn't exist.
ValueError: If scale is not 2, 3, or 4.
"""
if scale not in (2, 3, 4):
raise ValueError(f"scale must be 2, 3, or 4; got {scale}")
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)
_upsampler = upsampler
close_upsampler = upsampler is None
if _upsampler is None:
_upsampler = _build_upsampler()
img = Image.open(input_file).convert("RGB")
img_np = np.array(img)
t0 = time.perf_counter()
result, _ = _upsampler.enhance(img_np, outscale=scale)
t_elapsed = round(time.perf_counter() - t0, 3)
result_img = Image.fromarray(result)
result_img.save(output_file)
w, h = img.size
rw, rh = result_img.size
print(f"Saved {output_file} ({w}×{h} → {rw}×{rh}, {t_elapsed}s)")
if close_upsampler:
del _upsampler
def upscales(
input_paths: Sequence[str],
output_dir: str,
scale: int = 2,
upsampler: RealESRGANer | None = None,
) -> None:
"""Upscale multiple images, reusing one model session.
Output files are named like ``{stem}_x{scale}.png`` in ``output_dir``.
Args:
input_paths: List of input image paths.
output_dir: Directory for output PNGs.
scale: Output scale factor (2, 3, or 4).
upsampler: A pre-built RealESRGANer. If None, one is created and
discarded after the batch.
"""
if scale not in (2, 3, 4):
raise ValueError(f"scale must be 2, 3, or 4; got {scale}")
out_dir = Path(output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
_upsampler = upsampler
close_upsampler = upsampler is None
if _upsampler is None:
_upsampler = _build_upsampler()
total = 0.0
for input_path in input_paths:
stem = Path(input_path).stem
output_path = out_dir / f"{stem}_x{scale}.png"
img = Image.open(input_path).convert("RGB")
img_np = np.array(img)
t0 = time.perf_counter()
result, _ = _upsampler.enhance(img_np, outscale=scale)
t_elapsed = time.perf_counter() - t0
total += t_elapsed
result_img = Image.fromarray(result)
result_img.save(output_path)
w, h = img.size
rw, rh = result_img.size
print(f"Saved {output_path} ({w}×{h} → {rw}×{rh}, {t_elapsed:.3f}s)")
print(f"Done: {len(input_paths)} images, {total:.2f}s total")
if close_upsampler:
del _upsampler