backend/model.py

"""
Model download and CPU-optimized ORT session.
Uses the original FP32 .onnx weights as provided.
"""

import sys
import multiprocessing as mp
from pathlib import Path

import requests
import onnxruntime as ort

HF_URL = "https://huggingface.co/Subh775/Dis-Seg-Former/resolve/main/export/rfdetr-seg-nano.onnx"
MODEL_DIR = Path("/tmp/dis_seg_model")
MODEL_PATH = MODEL_DIR / "rfdetr-seg-nano.onnx"
DOWNLOAD_TIMEOUT = 300


def _download():
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    print(f"[model] Downloading {HF_URL} ...")
    r = requests.get(HF_URL, stream=True, timeout=DOWNLOAD_TIMEOUT)
    r.raise_for_status()
    with open(MODEL_PATH, "wb") as f:
        for chunk in r.iter_content(chunk_size=8192):
            f.write(chunk)
    print(f"[model] Downloaded ({MODEL_PATH.stat().st_size // 1024} KB).")


def load_model():
    try:
        if not MODEL_PATH.exists() or MODEL_PATH.stat().st_size == 0:
            _download()
        else:
            print(f"[model] Using cached weights at {MODEL_PATH}")
    except Exception as e:
        print(f"[model] FATAL: download failed — {e}", file=sys.stderr)
        sys.exit(1)

    print("[model] Creating ONNX Runtime session (CPU, optimized)...")
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    so.execution_mode = ort.ExecutionMode.ORT_PARALLEL
    cpu_count = mp.cpu_count()
    so.intra_op_num_threads = max(1, cpu_count)
    so.inter_op_num_threads = max(1, cpu_count // 2)
    so.enable_mem_pattern = True
    so.enable_cpu_mem_arena = True

    session = ort.InferenceSession(
        str(MODEL_PATH),
        sess_options=so,
        providers=["CPUExecutionProvider"],
    )

    inp = session.get_inputs()[0]
    print(f"[model] Input: name={inp.name}, shape={inp.shape}, threads(intra/inter)={so.intra_op_num_threads}/{so.inter_op_num_threads}")
    return session


def warmup(session):
    """Two warmup passes so first user request hits the optimized path."""
    import numpy as np
    inp = session.get_inputs()[0]
    _, _, h, w = inp.shape
    dummy = np.zeros((1, 3, h, w), dtype=np.float32)
    session.run(None, {inp.name: dummy})
    session.run(None, {inp.name: dummy})
    print("[model] Warm-up complete.")