| """WeMM-JEV text and vision inference.""" |
|
|
| from __future__ import annotations |
|
|
| import contextlib |
| import gc |
| import os |
| import sys |
| import threading |
| from pathlib import Path |
|
|
| from ._compat import Engine, Unsupported |
| from . import _images as vimg |
| from ._images import resize_policy |
|
|
| os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1") |
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
|
|
| _LOAD_LOCK = threading.Lock() |
|
|
|
|
| class D3Engine(Engine): |
| name = "d3" |
| latency = ( |
| "Device-synchronized in-process request wall time including prompt rendering and tokenization (and, " |
| "for image and video requests, their decoding and preprocessing); one request per call, its questions in " |
| "request order in batches of batch_size; excludes model loading." |
| ) |
|
|
| def __init__( |
| self, |
| model: str | None = None, |
| revision: str | None = None, |
| device: str | None = None, |
| batch_size: int = 8, |
| verify: str = "fast", |
| model_name: str | None = "WeMM-JEV", |
| permutation_average: str | bool = False, |
| max_pixels: int | None = vimg.MAX_PIXELS, |
| image_root: str | None = None, |
| replicas: int | None = None, |
| full_warmup: str | bool = False, |
| enable_videos: str | bool = False, |
| max_batch_tokens: int = 8192, |
| **options, |
| ): |
| if options: |
| raise TypeError(f"unknown engine options {sorted(options)}") |
| if not model: |
| model = str(Path(__file__).absolute().parent.parent) |
| flag = str(permutation_average).strip().lower() |
| if flag not in ("0", "1", "false", "true"): |
| raise ValueError("permutation_average must be 0 or 1") |
| averaged = flag in ("1", "true") |
| fw_flag = str(full_warmup).strip().lower() |
| if fw_flag not in ("0", "1", "false", "true"): |
| raise ValueError("full_warmup must be 0 or 1") |
| self.full_warmup = fw_flag in ("1", "true") |
| video_flag = str(enable_videos).strip().lower() |
| if video_flag not in ("0", "1", "false", "true"): |
| raise ValueError("enable_videos must be 0 or 1") |
| self.enable_videos = video_flag in ("1", "true") |
|
|
| import torch |
|
|
| self.torch = torch |
| self.model_id = str(model) |
| self.revision = revision |
| self.batch_size = int(batch_size) |
| self.max_batch_tokens = int(max_batch_tokens) |
| if self.batch_size < 1 or self.max_batch_tokens < 0: |
| raise ValueError("batch_size must be positive and max_batch_tokens nonnegative") |
| self.verify = verify |
| self.model_name = model_name |
| self.permutation_average = averaged |
| self.max_pixels = max_pixels |
| self.image_root = image_root |
|
|
| gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 |
| if device in (None, "cuda"): |
| primary_device = "cuda:0" if gpus > 0 else "cpu" |
| self._multi_gpu = gpus > 1 |
| self._num_gpus = max(1, gpus) |
| else: |
| primary_device = str(device) |
| self._multi_gpu = False |
| self._num_gpus = 1 |
|
|
| super().__init__( |
| model=self.model_id, |
| revision=revision, |
| device=primary_device, |
| batch_size=self.batch_size, |
| verify=verify, |
| model_name=model_name, |
| max_pixels=max_pixels, |
| enable_videos=self.enable_videos, |
| max_batch_tokens=self.max_batch_tokens, |
| **({"permutation_average": 1} if averaged else {}), |
| ) |
|
|
| hf_model, decision = self._load_replica(primary_device) |
| self.model = hf_model |
| self.decision = decision |
| self._slots = {0: (hf_model, decision, threading.Lock())} |
| self._local = threading.local() |
| self._next_slot = 0 |
| self._slots_lock = threading.Lock() |
|
|
| if replicas is not None: |
| n_rep = int(replicas) |
| if n_rep > 1: |
| if gpus < n_rep: |
| raise RuntimeError(f"replicas={n_rep} needs {n_rep} CUDA devices, found {gpus}") |
| self._multi_gpu = True |
| self._num_gpus = n_rep |
| for idx in range(1, n_rep): |
| m_i, d_i = self._load_replica(f"cuda:{idx}") |
| self._slots[idx] = (m_i, d_i, threading.Lock()) |
|
|
| self.provenance = { |
| **self.decision.provenance(), |
| "max_pixels": self.max_pixels, |
| "max_batch_tokens": self.max_batch_tokens, |
| "image_resize": resize_policy(self.max_pixels) + "; then the checkpoint's own processor", |
| } |
|
|
| def _load_replica(self, device_str: str): |
| with _LOAD_LOCK: |
| local = Path(os.fspath(self.model_id)).expanduser() |
| if local.is_dir(): |
| pkg_dir = os.path.abspath(os.fspath(local)) |
| if pkg_dir not in sys.path: |
| sys.path.insert(0, pkg_dir) |
| from transformers import AutoModel |
|
|
| if device_str.startswith("cuda"): |
| self.torch.cuda.set_device(device_str) |
| hf_model = AutoModel.from_pretrained( |
| self.model_id, |
| trust_remote_code=True, |
| revision=self.revision, |
| device=device_str, |
| batch_size=self.batch_size, |
| verify=self.verify, |
| ) |
| decision = getattr(hf_model, "runtime", None) |
| if decision is None: |
| raise ValueError( |
| f"{self.model_id!r} did not load as a d3 decision package (missing decision_config.json)" |
| ) |
| if self.model_name: |
| decision.model_name = self.model_name |
| if self.permutation_average: |
| decision.permutation_average = True |
| |
| |
| if not self.enable_videos and getattr(decision, "video_unavailable", None) is None: |
| decision.video_unavailable = "disabled for text/image evaluation (enable_videos=0)" |
| return hf_model, decision |
|
|
| def _acquire_slot(self): |
| slot = getattr(self._local, "slot", None) |
| if slot is None: |
| with self._slots_lock: |
| slot = self._next_slot if self._multi_gpu else 0 |
| self._next_slot = (self._next_slot + 1) % max(1, self._num_gpus) |
| self._local.slot = slot |
| entry = self._slots.get(slot) |
| if entry is None: |
| with self._slots_lock: |
| entry = self._slots.get(slot) |
| if entry is None: |
| hf_model, decision = self._load_replica(f"cuda:{slot}") |
| entry = (hf_model, decision, threading.Lock()) |
| self._slots[slot] = entry |
| return entry |
|
|
| @contextlib.contextmanager |
| def _device_ctx(self, decision): |
| dev = getattr(decision, "device", None) |
| if dev is not None and getattr(dev, "type", None) == "cuda": |
| with self.torch.cuda.device(dev): |
| yield |
| else: |
| yield |
|
|
| def warmup(self): |
| with self._device_ctx(self.decision): |
| super().warmup() |
| if self.full_warmup or getattr(self.decision, "fast", None) is not None: |
| self.warmup_seconds = self.decision.warmup() |
|
|
| def runtime(self): |
| return self.decision.runtime_info() |
|
|
| def synchronize(self): |
| for _, decision, _ in list(self._slots.values()): |
| decision.synchronize() |
|
|
| def _prepare_images(self, images, decision): |
| if not images: |
| return None, None |
| from PIL import Image |
|
|
| pil_images = [] |
| image_pixels = [] |
| for item in images: |
| if isinstance(item, Image.Image): |
| im = item.convert("RGB") |
| if self.max_pixels and im.width * im.height > self.max_pixels: |
| im = im.resize(vimg.fit(im.width, im.height, self.max_pixels), Image.LANCZOS) |
| else: |
| im = vimg.pil(item, self.image_root, self.max_pixels).convert("RGB") |
| pil_images.append(im) |
| image_pixels.append(im.width * im.height) |
| image_tokens = None |
| proc = getattr(decision, "processor", None) |
| if proc is not None and getattr(proc, "image_processor", None) is not None: |
| ip = proc.image_processor |
| merge = getattr(ip, "merge_size", 2) ** 2 |
| try: |
| image_tokens = sum( |
| ip.get_number_of_image_patches(im.height, im.width, {}) // merge |
| for im in pil_images |
| ) |
| except Exception: |
| image_tokens = None |
| raw = {"image_pixels": image_pixels} |
| if image_tokens is not None: |
| raw["image_tokens"] = image_tokens |
| return pil_images, raw |
|
|
| def __call__(self, state, questions, images=None, videos=None, prompt=None): |
| if videos and not self.enable_videos: |
| raise Unsupported("video evaluation requires --option enable_videos=1 and a working video decoder") |
| _, decision, lock = self._acquire_slot() |
| pil_images, raw = self._prepare_images(images, decision) |
| with lock, self._device_ctx(decision): |
| prepared = decision.prepare(state, questions, pil_images, videos) |
| over = [ |
| e["message"] |
| for e in prepared.errors.values() |
| if e["error"] == "max_length_exceeded" |
| ] |
| if over: |
| raise Unsupported(over[0]) |
| if prepared.errors: |
| raise ValueError( |
| "invalid questions: " |
| + "; ".join(f"{k}: {e['message']}" for k, e in prepared.errors.items()) |
| ) |
| probabilities, tokens = self._run_prepared(decision, prepared) |
| response = decision.respond(prepared, probabilities, tokens) |
| failed = { |
| k: a["message"] for k, a in response["answers"].items() if "error" in a |
| } |
| if failed: |
| raise ValueError(f"invalid model output: {failed}") |
| return response, raw |
|
|
| def _run_prepared(self, decision, prepared): |
| """Bound padded tokens per forward without truncating questions or context.""" |
| original = decision.batch_size |
| sizes = list(getattr(prepared, "lengths", {}).values()) |
| sizes.extend(len(s) for s in getattr(prepared, "sequences", {}).values()) |
| sizes.extend(len(s) for s in getattr(prepared, "reversed_sequences", {}).values()) |
| sizes.extend(getattr(prepared, "reversed_lengths", {}).values()) |
| batch = original |
| if self.max_batch_tokens and sizes: |
| width = max(sizes) * (2 if self.permutation_average else 1) |
| batch = min(batch, max(1, self.max_batch_tokens // max(1, width))) |
| try: |
| while True: |
| decision.batch_size = batch |
| try: |
| return decision.run(prepared) |
| except self.torch.OutOfMemoryError: |
| if batch == 1: |
| raise |
| |
| batch = max(1, batch // 2) |
| gc.collect() |
| self.torch.cuda.empty_cache() |
| finally: |
| decision.batch_size = original |
|
|
| def close(self): |
| with self._slots_lock: |
| self._slots.clear() |
| self.model = None |
| self.decision = None |
| if self.torch.cuda.is_available(): |
| self.torch.cuda.empty_cache() |
|
|
