From 8b3431ff12ef18abb20148b3af9081191787309b Mon Sep 17 00:00:00 2001
From: Michele Rossi <mikymod@gmail.com>
Date: Thu, 9 Jul 2026 20:54:56 +0200
Subject: [PATCH] Fix bg removal in final assets

---
 vnassets/session.py |  4 ++--
 vnassets/upscale.py | 53 +++++++++++++++++++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/vnassets/session.py b/vnassets/session.py
index fc8c7b9..2378e07 100644
--- a/vnassets/session.py
+++ b/vnassets/session.py
@@ -407,10 +407,10 @@ class VnAssetsSession:
 
         The upsampler is lazy-loaded on first call and reused across
         calls. Uses ``RealESRGAN_x4plus_anime_6B`` (17 MB,
-        auto-downloaded).
+        auto-downloaded). Preserves alpha channel when present.
 
         Args:
-            input_path: Path to the input RGB image.
+            input_path: Path to the input image (RGB or RGBA).
             output_path: Where to save the upscaled PNG.
             scale: Output scale factor (2, 3, or 4).
 
diff --git a/vnassets/upscale.py b/vnassets/upscale.py
index 9f8cc7b..129e926 100644
--- a/vnassets/upscale.py
+++ b/vnassets/upscale.py
@@ -72,13 +72,32 @@ def _build_upsampler(device: str = "cuda") -> RealESRGANer:
     return upsampler
 
 
+def _load_rgb_alpha(img: Image.Image) -> tuple[np.ndarray, np.ndarray | None]:
+    """Extract RGB array and optional alpha channel from a PIL image.
+
+    Returns (rgb_np, alpha_np). ``alpha_np`` is None when the image has no
+    alpha channel. For paletted images, the palette is expanded first.
+    """
+    if img.mode == "P":
+        img = img.convert("RGBA" if "transparency" in img.info else "RGB")
+    if img.mode in ("RGBA", "LA"):
+        rgb = img.convert("RGB")
+        alpha = img.getchannel("A")
+        return np.array(rgb), np.array(alpha)
+    return np.array(img.convert("RGB")), None
+
+
 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.
+    """Upscale a single image. Preserves alpha channel when present.
+
+    When the input is RGBA, the RGB channels are upscaled through the
+    model and the alpha channel is resized with LANCZOS interpolation.
+    Output format matches input: RGB for RGB inputs, RGBA for RGBA inputs.
 
     Args:
         input_path: Path to the input image.
@@ -107,17 +126,25 @@ def upscale(
     if _upsampler is None:
         _upsampler = _build_upsampler()
 
-    img = Image.open(input_file).convert("RGB")
-    img_np = np.array(img)
+    src = Image.open(input_file)
+    rgb_np, alpha_np = _load_rgb_alpha(src)
+    w, h = src.size
+    src.close()
 
     t0 = time.perf_counter()
-    result, _ = _upsampler.enhance(img_np, outscale=scale)
+    result, _ = _upsampler.enhance(rgb_np, outscale=scale)
     t_elapsed = round(time.perf_counter() - t0, 3)
 
     result_img = Image.fromarray(result)
+
+    if alpha_np is not None:
+        alpha_img = Image.fromarray(alpha_np)
+        alpha_scaled = alpha_img.resize(result_img.size, Image.LANCZOS)
+        result_img = result_img.convert("RGBA")
+        result_img.putalpha(alpha_scaled)
+
     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)")
 
@@ -158,17 +185,25 @@ def upscales(
         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)
+        src = Image.open(input_path)
+        rgb_np, alpha_np = _load_rgb_alpha(src)
+        w, h = src.size
+        src.close()
 
         t0 = time.perf_counter()
-        result, _ = _upsampler.enhance(img_np, outscale=scale)
+        result, _ = _upsampler.enhance(rgb_np, outscale=scale)
         t_elapsed = time.perf_counter() - t0
         total += t_elapsed
 
         result_img = Image.fromarray(result)
+
+        if alpha_np is not None:
+            alpha_img = Image.fromarray(alpha_np)
+            alpha_scaled = alpha_img.resize(result_img.size, Image.LANCZOS)
+            result_img = result_img.convert("RGBA")
+            result_img.putalpha(alpha_scaled)
+
         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)")