add Real-ESRGAN upscaling (2x/3x/4x) with anime model
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
92
README.md
92
README.md
@@ -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 |
|
||||
|
||||
Reference in New Issue
Block a user