Upload export_birefnet_onnx.py

export_birefnet_onnx.py

@@ -1,30 +1,25 @@
 #!/usr/bin/env python3
 """
-Export BiRefNet .pth weights to ONNX.
-
-Works with the environment used by BiRefNet demo-style setups:
+BiRefNet (.pth) -> ONNX exporter that works with:
 - Python 3.10
-- torch==2.0.1+cu118
-- transformers==4.42.4  (IMPORTANT: newer transformers may require torch>=2.1 and will disable torch)
-
-Example:
-  python export_birefnet_onnx.py \
-    --code_dir birefnet_code \
-    --weights weights/birefnet_finetuned_toonout.pth \
-    --output birefnet_finetuned_toonout.onnx \
-    --img_size 1024 \
-    --opset 17 \
-    --device cuda \
-    --external_data
+- torch==2.0.1 (+cu118 recommended)
+- transformers==4.42.4
+
+Fixes:
+- BiRefNet HF code uses relative imports (e.g. from .BiRefNet_config import ...),
+  so --code_dir must be imported as a *package*.
+- Some public scripts pass use_external_data_format to torch.onnx.export, but
+  torch 2.0.1 does NOT support that keyword.
+- Some checkpoints are saved from torch.compile and have keys prefixed with `_orig_mod.`.
 """
 
 from __future__ import annotations
 
 import argparse
+import importlib
 import os
 import sys
-from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Iterable
 
 import torch
 
@@ -34,200 +29,266 @@ def _print_env(device: str) -> None:
     print("Python:", sys.version.replace("\n", " "))
     print("Torch:", torch.__version__)
     print("CUDA available:", torch.cuda.is_available())
-    if device.startswith("cuda") and torch.cuda.is_available():
-        idx = 0
+    if torch.cuda.is_available():
         try:
             idx = torch.cuda.current_device()
+            print("CUDA device:", torch.cuda.get_device_name(idx))
         except Exception:
             pass
-        try:
-            name = torch.cuda.get_device_name(idx)
-        except Exception:
-            name = "cuda"
-        print("CUDA device:", name)
+    print("Requested device:", device)
 
 
-def _ensure_transformers_torch_backend_ok() -> None:
-    # If transformers is installed but thinks torch is unavailable, BiRefNet will import dummy torch objects.
+def _try_register_deform_conv2d() -> bool:
+    """
+    Optional: register ONNX symbolic for torchvision's DeformConv2d.
+    Provided by deform-conv2d-onnx-exporter.
+    """
     try:
-        import transformers  # noqa
-        from transformers.utils import is_torch_available  # noqa
-
-        if not is_torch_available():
-            raise RuntimeError(
-                "Transformers is installed but has DISABLED the PyTorch backend.\n"
-                "This usually happens when your transformers version requires a newer torch.\n\n"
-                "Fix (recommended):\n"
-                "  pip uninstall -y transformers tokenizers\n"
-                "  pip install transformers==4.42.4 huggingface_hub==0.23.4\n"
-            )
-    except ModuleNotFoundError:
-        # BiRefNet HF-style code requires transformers; we'll let the import fail later with a clear error.
-        pass
+        import deform_conv2d_onnx_exporter  # type: ignore
 
-
-def _try_register_deformconv_exporter() -> None:
-    ok = False
-    try:
-        import deform_conv2d_onnx_exporter as d  # type: ignore
-
-        # Try common entry points
-        for name in (
-            "register_deform_conv2d_onnx_exporter",
-            "register",
-            "setup",
-        ):
-            fn = getattr(d, name, None)
-            if callable(fn):
-                fn()
-                ok = True
+        deform_conv2d_onnx_exporter.register_deform_conv2d_onnx_op()
+        print("DeformConv2d ONNX exporter: OK")
+        return True
+    except Exception as e:
+        print("DeformConv2d ONNX exporter: NOT LOADED (may fail if model uses DeformConv)")
+        print("  Reason:", repr(e))
+        return False
+
+
+def _ensure_pkg_and_import(code_dir: str):
+    """
+    Make sure code_dir is a real python package, then import <pkg>.birefnet
+    so that relative imports inside birefnet.py work.
+    """
+    code_dir = os.path.abspath(code_dir)
+    if not os.path.isdir(code_dir):
+        raise FileNotFoundError(f"--code_dir not found or not a directory: {code_dir}")
+
+    init_py = os.path.join(code_dir, "__init__.py")
+    if not os.path.exists(init_py):
+        # create empty __init__.py to make it a package
+        open(init_py, "a", encoding="utf-8").close()
+
+    pkg_name = os.path.basename(code_dir.rstrip("/"))
+    parent_dir = os.path.dirname(code_dir)
+    if parent_dir not in sys.path:
+        sys.path.insert(0, parent_dir)
+
+    # Import as package to satisfy relative imports
+    mod = importlib.import_module(f"{pkg_name}.birefnet")
+    return mod, pkg_name
+
+
+def _extract_state_dict(ckpt_obj: Any) -> Dict[str, torch.Tensor]:
+    """
+    Accepts various checkpoint formats and returns a plain state_dict.
+    """
+    if isinstance(ckpt_obj, dict):
+        # common nesting keys
+        for k in ("state_dict", "model_state_dict", "model", "net", "params", "ema"):
+            v = ckpt_obj.get(k, None)
+            if isinstance(v, dict):
+                ckpt_obj = v
                 break
-    except Exception:
-        ok = False
 
-    print("DeformConv2d ONNX exporter:", "OK" if ok else "NOT LOADED (may fail if model uses DeformConv)")
+    if not isinstance(ckpt_obj, dict):
+        raise RuntimeError("Unsupported checkpoint format: expected a dict/state_dict.")
 
+    # At this point it should be {str: Tensor}
+    sd: Dict[str, torch.Tensor] = {}
+    for k, v in ckpt_obj.items():
+        if isinstance(k, str) and torch.is_tensor(v):
+            sd[k] = v
 
-def _load_state_dict(weights_path: Path) -> Dict[str, Any]:
-    sd = torch.load(str(weights_path), map_location="cpu")
-    if isinstance(sd, dict) and "state_dict" in sd and isinstance(sd["state_dict"], dict):
-        sd = sd["state_dict"]
+    if not sd:
+        raise RuntimeError("Checkpoint dict contained no tensor parameters.")
+    return sd
 
-    if not isinstance(sd, dict):
-        raise ValueError("Weights file did not contain a state_dict-like dict.")
 
-    # Remove 'module.' prefixes from DDP-trained checkpoints
-    clean = {}
+def _normalize_state_dict_keys(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """
+    Fix common prefixes:
+    - torch.compile checkpoints: `_orig_mod.`
+    - DataParallel / DDP: `module.`
+    """
+    out: Dict[str, torch.Tensor] = {}
     for k, v in sd.items():
-        nk = k[7:] if k.startswith("module.") else k
-        clean[nk] = v
-    return clean
-
-
-class _OutputWrapper(torch.nn.Module):
-    def __init__(self, model: torch.nn.Module):
+        nk = k
+        if nk.startswith("_orig_mod."):
+            nk = nk[len("_orig_mod.") :]
+        if nk.startswith("module."):
+            nk = nk[len("module.") :]
+        out[nk] = v
+    return out
+
+
+def _iter_tensors(x: Any) -> Iterable[torch.Tensor]:
+    if torch.is_tensor(x):
+        yield x
+    elif isinstance(x, dict):
+        for v in x.values():
+            yield from _iter_tensors(v)
+    elif isinstance(x, (list, tuple)):
+        for v in x:
+            yield from _iter_tensors(v)
+
+
+def _pick_best_output(out: Any, img_size: int | None) -> torch.Tensor:
+    """
+    BiRefNet forward can return nested structures (list/tuple/dict).
+    We want a single mask tensor [N,1,H,W] if possible.
+    """
+    tensors = list(_iter_tensors(out))
+    if not tensors:
+        raise RuntimeError("Model forward produced no tensors to export.")
+
+    # Prefer rank-4 tensors
+    cands = [t for t in tensors if t.dim() == 4]
+
+    # Prefer exact H/W match if provided
+    if img_size is not None and cands:
+        cands_hw = [t for t in cands if int(t.shape[-2]) == img_size and int(t.shape[-1]) == img_size]
+        if cands_hw:
+            cands = cands_hw
+
+    # Prefer single-channel outputs
+    if cands:
+        cands_c1 = [t for t in cands if int(t.shape[1]) == 1]
+        if cands_c1:
+            cands = cands_c1
+
+    return cands[0] if cands else tensors[0]
+
+
+class _ExportWrapper(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, img_size: int | None):
         super().__init__()
         self.model = model
+        self.img_size = img_size
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = self.model(x)
-
-        # BiRefNet sometimes returns lists/tuples/dicts; ONNX export wants a tensor.
-        if isinstance(y, torch.Tensor):
-            return y
-
-        if isinstance(y, (list, tuple)) and len(y) > 0:
-            # Most BiRefNet variants put the final prediction at the end
-            last = y[-1]
-            if isinstance(last, torch.Tensor):
-                return last
-            # fallback: first tensor found
-            for item in y:
-                if isinstance(item, torch.Tensor):
-                    return item
-
-        if isinstance(y, dict):
-            for key in ("pred", "mask", "out", "logits"):
-                if key in y and isinstance(y[key], torch.Tensor):
-                    return y[key]
-            # fallback: first tensor value
-            for v in y.values():
-                if isinstance(v, torch.Tensor):
-                    return v
-
-        raise TypeError(f"Model forward returned unsupported type for ONNX export: {type(y)}")
+        out = self.model(x)
+        y = _pick_best_output(out, self.img_size)
+        return y
 
 
 def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--code_dir", type=str, required=True, help="Folder containing birefnet.py (downloaded BiRefNet code)")
-    ap.add_argument("--weights", type=str, required=True, help="Path to .pth weights")
-    ap.add_argument("--output", type=str, required=True, help="Output .onnx path")
-    ap.add_argument("--img_size", type=int, default=1024, help="Square input size (default 1024)")
-    ap.add_argument("--opset", type=int, default=17, help="ONNX opset (default 17)")
-    ap.add_argument("--device", type=str, default="cuda", help="cuda or cpu")
-    ap.add_argument("--half", action="store_true", help="Export in fp16 (not always supported)")
-    ap.add_argument("--dynamic_axes", action="store_true", help="Enable dynamic batch axis")
-    ap.add_argument("--external_data", action="store_true", help="Use external data format (for >2GB models)")
-    args = ap.parse_args()
-
-    code_dir = Path(args.code_dir).resolve()
-    weights_path = Path(args.weights).resolve()
-    out_path = Path(args.output).resolve()
+    p = argparse.ArgumentParser()
+    p.add_argument("--code_dir", required=True, help="Folder that contains birefnet.py and BiRefNet_config.py")
+    p.add_argument("--weights", required=True, help="Path to .pth weights")
+    p.add_argument("--output", required=True, help="Output ONNX path, e.g. out.onnx")
+    p.add_argument("--img_size", type=int, default=1024, help="Dummy input resolution (square), default 1024")
+    p.add_argument("--opset", type=int, default=17, help="ONNX opset, default 17")
+    p.add_argument("--device", default="cuda", choices=["cuda", "cpu"], help="cuda or cpu")
+    p.add_argument("--dynamic", action="store_true", help="Export dynamic H/W axes (may break export)")
+    p.add_argument(
+        "--external_data",
+        action="store_true",
+        help="After export, re-save ONNX using external data (.onnx + .onnx.data).",
+    )
+    p.add_argument("--skip_onnx_check", action="store_true", help="Skip onnx.checker.check_model()")
+    args = p.parse_args()
 
     _print_env(args.device)
-    _ensure_transformers_torch_backend_ok()
-    _try_register_deformconv_exporter()
+    _try_register_deform_conv2d()
 
-    if not code_dir.exists():
-        raise FileNotFoundError(f"--code_dir not found: {code_dir}")
-    if not weights_path.exists():
-        raise FileNotFoundError(f"--weights not found: {weights_path}")
-
-    sys.path.insert(0, str(code_dir))
-    try:
-        from birefnet import BiRefNet  # type: ignore
-    except Exception as e:
-        raise RuntimeError(
-            "Failed to import BiRefNet from your --code_dir.\n"
-            f"code_dir={code_dir}\n"
-            "Make sure birefnet.py exists there.\n"
-            f"Original error: {e}"
-        )
+    # Import model properly (as a package)
+    birefnet_mod, pkg_name = _ensure_pkg_and_import(args.code_dir)
+    if not hasattr(birefnet_mod, "BiRefNet"):
+        raise RuntimeError(f"BiRefNet class not found in {pkg_name}.birefnet")
+    BiRefNet = getattr(birefnet_mod, "BiRefNet")
 
     print("== Building model ==")
     model = BiRefNet(bb_pretrained=False)
+    model.eval()
 
     print("== Loading weights ==")
-    sd = _load_state_dict(weights_path)
-    missing, unexpected = model.load_state_dict(sd, strict=False)
-    print(f"Missing keys: {len(missing)}; Unexpected keys: {len(unexpected)}")
-    if len(missing) > 0:
-        print("  (missing example):", missing[:10])
-    if len(unexpected) > 0:
-        print("  (unexpected example):", unexpected[:10])
-
-    device = torch.device(args.device if args.device != "cuda" else ("cuda" if torch.cuda.is_available() else "cpu"))
-    model.eval().to(device)
-
-    dtype = torch.float16 if args.half else torch.float32
-    wrapper = _OutputWrapper(model).to(device)
-
-    dummy = torch.randn(1, 3, args.img_size, args.img_size, device=device, dtype=dtype)
-
-    # Quick forward to ensure it runs
+    ckpt = torch.load(args.weights, map_location="cpu")
+    sd = _extract_state_dict(ckpt)
+    sd = _normalize_state_dict_keys(sd)
+
+    incompatible = model.load_state_dict(sd, strict=False)
+    missing = list(getattr(incompatible, "missing_keys", []))
+    unexpected = list(getattr(incompatible, "unexpected_keys", []))
+    print(f"Loaded state_dict. Missing keys: {len(missing)}  Unexpected keys: {len(unexpected)}")
+    if missing:
+        print("  (first 20 missing):", missing[:20])
+    if unexpected:
+        print("  (first 20 unexpected):", unexpected[:20])
+
+    if args.device == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("You asked for --device cuda but CUDA is not available.")
+        model = model.to("cuda")
+        dev = "cuda"
+    else:
+        model = model.to("cpu")
+        dev = "cpu"
+
+    wrapper = _ExportWrapper(model, img_size=args.img_size)
+    wrapper.eval()
+
+    dummy = torch.randn(1, 3, args.img_size, args.img_size, device=dev)
+
+    print("== Forward probe ==")
     with torch.no_grad():
-        _ = wrapper(dummy)
-
-    out_path.parent.mkdir(parents=True, exist_ok=True)
+        probe_out = wrapper(dummy)
+    print("Picked output tensor shape:", tuple(probe_out.shape), "dtype:", probe_out.dtype)
 
     print("== Exporting ONNX ==")
-    input_names = ["input"]
-    output_names = ["output"]
+    out_path = os.path.abspath(args.output)
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
 
+    input_names = ["input"]
+    output_names = ["mask"]
     dynamic_axes = None
-    if args.dynamic_axes:
-        dynamic_axes = {"input": {0: "batch"}, "output": {0: "batch"}}
-
-    # External data format is needed if model > 2GB.
-    use_external_data_format = bool(args.external_data)
-
-    torch.onnx.export(
-        wrapper,
-        dummy,
-        str(out_path),
-        export_params=True,
-        opset_version=args.opset,
-        do_constant_folding=True,
-        input_names=input_names,
-        output_names=output_names,
-        dynamic_axes=dynamic_axes,
-        use_external_data_format=use_external_data_format,
-    )
+    if args.dynamic:
+        dynamic_axes = {
+            "input": {0: "batch", 2: "height", 3: "width"},
+            "mask": {0: "batch", 2: "height", 3: "width"},
+        }
+
+    with torch.no_grad():
+        # IMPORTANT: torch 2.0.1 does NOT support use_external_data_format.
+        torch.onnx.export(
+            wrapper,
+            dummy,
+            out_path,
+            export_params=True,
+            opset_version=args.opset,
+            do_constant_folding=True,
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+        )
 
-    print("Saved:", out_path)
-    if use_external_data_format:
-        print("External data format enabled: you may also have an .onnx.data file next to the ONNX.")
+    print("Output:", out_path)
+
+    if args.external_data or (not args.skip_onnx_check):
+        import onnx  # type: ignore
+
+        print("== Loading ONNX ==")
+        onnx_model = onnx.load(out_path)
+
+        if not args.skip_onnx_check:
+            print("== ONNX checker ==")
+            onnx.checker.check_model(onnx_model)
+            print("ONNX checker: OK")
+
+        if args.external_data:
+            print("== Saving external data ==")
+            data_name = os.path.basename(out_path) + ".data"
+            onnx.save_model(
+                onnx_model,
+                out_path,
+                save_as_external_data=True,
+                all_tensors_to_one_file=True,
+                location=data_name,
+                size_threshold=1024,  # bytes; moves almost everything out
+            )
+            print("Saved external-data ONNX:")
+            print("  Model:", out_path)
+            print("  Data :", os.path.join(os.path.dirname(out_path), data_name))
 
     print("== Done ==")