Delete raven/foundation_runtime.py

raven/foundation_runtime.py

@@ -1,179 +0,0 @@
-from __future__ import annotations
-
-import gc
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-from .foundation import FoundationEncoder
-
-
-class _FoundationWrapper(nn.Module):
-    def __init__(self, enc: FoundationEncoder):
-        super().__init__()
-        self.enc = enc
-
-    def forward(self, x):
-        return self.enc(x)
-
-
-class FoundationRuntime:
-
-
-    def __init__(self, cfg: dict, encoder: FoundationEncoder):
-        self.cfg = cfg
-        self.encoder = encoder
-        self.backend = "torch"
-        self.batch_size = max(1, int(cfg.get("torch_batch_size", 8)))
-        self.runtime_cache = Path(cfg.get("runtime_cache_dir", "cache/foundation_runtime"))
-        self.runtime_cache.mkdir(parents=True, exist_ok=True)
-        self._session: Any = None
-        self._compiled: Any = None
-
-        requested = str(cfg.get("accelerator", "auto")).strip().lower()
-        order = [requested] if requested not in {"auto", ""} else ["directml", "openvino", "torch"]
-
-        errors = []
-        for candidate in order:
-            try:
-                if candidate == "directml":
-                    self._init_directml()
-                elif candidate == "openvino":
-                    self._init_openvino()
-                elif candidate == "torch":
-                    self._init_torch()
-                else:
-                    raise ValueError(f"Unknown foundation accelerator: {candidate}")
-                print(f"[foundation-runtime] {self.backend}, batch={self.batch_size}")
-                break
-            except Exception as exc:
-                errors.append(f"{candidate}: {exc}")
-                self._session = None
-                self._compiled = None
-        else:
-            raise RuntimeError("No foundation runtime could be initialized: " + " | ".join(errors))
-
-    @property
-    def model_backend(self) -> str:
-        return self.encoder.backend
-
-    @property
-    def dim(self) -> int:
-        return int(self.encoder.dim)
-
-    def _fixed_batch_onnx_path(self, batch: int) -> Path:
-        safe = self.encoder.backend.replace("/", "_")
-        return self.runtime_cache / f"{safe}_b{batch}_fixed.onnx"
-
-    def _init_directml(self):
-        import onnxruntime as ort
-
-        available = ort.get_available_providers()
-        if "DmlExecutionProvider" not in available:
-            raise RuntimeError(f"DmlExecutionProvider unavailable; providers={available}")
-
-        batch = max(1, int(self.cfg.get("directml_batch_size", 4)))
-        onnx_path = self._fixed_batch_onnx_path(batch)
-        if not onnx_path.exists():
-            print(f"[foundation-runtime] exporting fixed-batch ONNX once -> {onnx_path}")
-            wrapper = _FoundationWrapper(self.encoder).eval()
-            dummy = torch.randn(batch, 3, 224, 224, dtype=torch.float32)
-            torch.onnx.export(
-                wrapper,
-                dummy,
-                str(onnx_path),
-                input_names=["image"],
-                output_names=["embedding"],
-                opset_version=18,
-                do_constant_folding=True,
-            )
-
-        opts = ort.SessionOptions()
-
-        opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
-        opts.enable_mem_pattern = False
-        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-        self._session = ort.InferenceSession(
-            str(onnx_path),
-            sess_options=opts,
-            providers=["DmlExecutionProvider", "CPUExecutionProvider"],
-        )
-        self.backend = "directml"
-        self.batch_size = batch
-
-    def _init_openvino(self):
-        import openvino as ov
-
-        batch = max(1, int(self.cfg.get("openvino_batch_size", 16)))
-        safe = self.encoder.backend.replace("/", "_")
-        xml = self.runtime_cache / f"{safe}_b{batch}.xml"
-
-        if not xml.exists():
-            print(f"[foundation-runtime] converting to OpenVINO IR once -> {xml}")
-            wrapper = _FoundationWrapper(self.encoder).eval()
-            dummy = torch.randn(batch, 3, 224, 224, dtype=torch.float32)
-            ov_model = ov.convert_model(wrapper, example_input=dummy)
-            ov.save_model(
-                ov_model,
-                str(xml),
-                compress_to_fp16=bool(self.cfg.get("openvino_compress_fp16", True)),
-            )
-
-        core = ov.Core()
-        config = {}
-
-        config["PERFORMANCE_HINT"] = "THROUGHPUT"
-        try:
-            config["INFERENCE_NUM_THREADS"] = int(self.cfg.get("openvino_threads", 6))
-        except Exception:
-            pass
-        try:
-            self._compiled = core.compile_model(str(xml), "CPU", config)
-        except Exception:
-
-            self._compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "THROUGHPUT"})
-        self.backend = "openvino"
-        self.batch_size = batch
-
-    def _init_torch(self):
-        self.encoder.eval().to("cpu")
-        self.backend = "torch"
-        self.batch_size = max(1, int(self.cfg.get("torch_batch_size", 8)))
-
-    def _pad(self, x: np.ndarray | torch.Tensor) -> tuple[np.ndarray, int]:
-        if isinstance(x, torch.Tensor):
-            a = x.detach().cpu().numpy().astype(np.float32, copy=False)
-        else:
-            a = np.asarray(x, dtype=np.float32)
-        n = int(a.shape[0])
-        if n == self.batch_size:
-            return a, n
-        if n > self.batch_size:
-            raise ValueError(f"Runtime batch {n} exceeds configured fixed batch {self.batch_size}")
-        pad = np.zeros((self.batch_size,) + tuple(a.shape[1:]), dtype=np.float32)
-        pad[:n] = a
-        return pad, n
-
-    @torch.inference_mode()
-    def infer(self, x: torch.Tensor) -> np.ndarray:
-        if self.backend == "torch":
-            return self.encoder(x.float()).cpu().numpy().astype(np.float32)
-
-        a, n = self._pad(x)
-        if self.backend == "directml":
-            y = self._session.run(["embedding"], {"image": a})[0]
-        elif self.backend == "openvino":
-            result = self._compiled([a])
-            y = result[self._compiled.output(0)]
-        else:
-            raise RuntimeError(self.backend)
-        return np.asarray(y[:n], dtype=np.float32)
-
-    def release_accelerator_memory(self):
-
-        self._session = None
-        self._compiled = None
-        gc.collect()