models.py

"""
models.py
4 background-removal models:
  fast    -> U2-Net        (rembg, MIT)
  quality -> BiRefNet-lite (transformers, MIT)
  hair    -> BEN2          (ben2, MIT) -- best for hair/portraits/fur
  best    -> BRIA RMBG-2.0 (transformers, non-commercial)
"""

import io
import logging
import warnings
import numpy as np
from PIL import Image

logger = logging.getLogger("hd_remover.models")

_cache: dict = {}


# ── U2-Net (fast) ─────────────────────────────────────────────────────────────
def load_u2net():
    if "u2net" in _cache:
        return _cache["u2net"]
    try:
        from rembg import new_session
        _cache["u2net"] = new_session("u2net")
        logger.info("U2-Net ready")
        return _cache["u2net"]
    except Exception as e:
        logger.error(f"U2-Net load failed: {e}")
        return None


def infer_fast(img: Image.Image) -> Image.Image:
    from rembg import remove
    if not load_u2net():
        raise RuntimeError("U2-Net not available")
    buf = io.BytesIO()
    img.save(buf, "PNG")
    buf.seek(0)
    return Image.open(io.BytesIO(remove(buf.read(), session=_cache["u2net"]))).convert("RGBA")


# ── BiRefNet-lite (quality) ────────────────────────────────────────────────────
def load_birefnet():
    if "birefnet" in _cache:
        return _cache["birefnet"]
    try:
        from transformers import AutoModelForImageSegmentation
        from torchvision import transforms
        m = AutoModelForImageSegmentation.from_pretrained(
            "ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
        m = m.float()
        m.eval()
        _cache["birefnet"] = m
        _cache["birefnet_tf"] = transforms.Compose([
            transforms.Resize((1024, 1024)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        logger.info("BiRefNet-lite ready")
        return m
    except Exception as e:
        logger.error(f"BiRefNet-lite load failed: {e}")
        return None


def infer_quality(img: Image.Image) -> Image.Image:
    if not load_birefnet():
        raise RuntimeError("BiRefNet-lite not available")
    return _seg("birefnet", img)


# ── BEN2 (hair) — MIT license, best for portraits/hair/fur ───────────────────
def load_ben2():
    if "ben2" in _cache:
        return _cache["ben2"]
    try:
        from ben2 import BEN_Base
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message=".*cuda.*CUDA is not available.*")
            m = BEN_Base.from_pretrained("PramaLLC/BEN2")
            m.eval()
        _cache["ben2"] = m
        logger.info("BEN2 ready")
        return m
    except Exception as e:
        logger.error(f"BEN2 load failed: {e}")
        return None


def infer_hair(img: Image.Image) -> Image.Image:
    """BEN2 — MIT license. Best for hair, portraits, fur edges."""
    if not load_ben2():
        raise RuntimeError("BEN2 not available")
    import torch
    model = _cache["ben2"]
    model.to(torch.device("cpu"))
    foreground = model.inference(img.convert("RGB"), refine_foreground=True)
    return foreground.convert("RGBA") if foreground.mode != "RGBA" else foreground


# ── BRIA RMBG-2.0 (best) — NON-COMMERCIAL ────────────────────────────────────
def load_rmbg(hf_token=None):
    if "rmbg" in _cache:
        return _cache["rmbg"]
    try:
        from transformers import AutoModelForImageSegmentation
        from torchvision import transforms
        m = AutoModelForImageSegmentation.from_pretrained(
            "briaai/RMBG-2.0", trust_remote_code=True, token=hf_token or None)
        m = m.float()
        m.eval()
        _cache["rmbg"] = m
        _cache["rmbg_tf"] = transforms.Compose([
            transforms.Resize((1024, 1024)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        logger.info("RMBG-2.0 ready")
        return m
    except Exception as e:
        logger.error(f"RMBG-2.0 load failed: {e}")
        return None


def infer_best(img: Image.Image, hf_token=None) -> Image.Image:
    if not load_rmbg(hf_token):
        raise RuntimeError("BRIA RMBG-2.0 not available — check HF_TOKEN")
    return _seg("rmbg", img)


# ── Shared segmentation helper (BiRefNet + RMBG) ─────────────────────────────
def _seg(key: str, img: Image.Image) -> Image.Image:
    import torch
    m = _cache[key]
    tf = _cache[key + "_tf"]
    orig = img.size
    inp = tf(img.convert("RGB")).unsqueeze(0).float()
    with torch.no_grad():
        out = m(inp)
        pred = (out[-1] if isinstance(out, (list, tuple)) else out).sigmoid()
    mask = Image.fromarray(
        (pred[0].squeeze().cpu().numpy() * 255).astype(np.uint8)
    ).resize(orig, Image.LANCZOS)
    r = img.convert("RGBA")
    r.putalpha(mask)
    return r


# ── Model registry ────────────────────────────────────────────────────────────
MODEL_INFO = {
    "fast":    {"name": "U2-Net",        "desc": "Fastest (~1-5s) — general use",               "license": "MIT"},
    "quality": {"name": "BiRefNet-lite", "desc": "High quality (~5-8s CPU) — products/objects", "license": "MIT"},
    "hair":    {"name": "BEN2",          "desc": "Best for hair & portraits (~10-15s CPU)",      "license": "MIT"},
    "best":    {"name": "BRIA RMBG-2.0", "desc": "Best quality (~20s CPU) — non-commercial",    "license": "Non-commercial"},
}


def preload_default():
    """Called at startup — loads U2-Net blocking, rest in background thread."""
    logger.info("Preloading U2-Net (startup)...")
    load_u2net()


def preload_all_bg():
    """Call in a background thread after startup."""
    logger.info("Background preload: BiRefNet-lite...")
    load_birefnet()
    logger.info("Background preload: BEN2...")
    load_ben2()
    logger.info("Background preload: RMBG-2.0...")
    load_rmbg()
    logger.info("All 4 models ready")


def models_ready() -> list:
    """Return list of currently loaded model keys."""
    return [k for k in _cache if not k.endswith("_tf")]