simit 0.1.0__py3-none-any.whl

simit/backends/hf.py

@@ -0,0 +1,336 @@
+"""Any Hugging Face ``transformers`` vision-language model.
+
+Requests are batched by a worker thread: whatever is queued when the GPU
+becomes free runs as one padded ``generate`` call (grouped by sampling
+settings), so concurrent pipeline steps share forward passes.
+"""
+
+from __future__ import annotations
+
+import inspect
+import queue
+import threading
+import traceback
+from concurrent.futures import Future
+from typing import Optional
+
+import torch
+from PIL import Image
+
+from ..utils import strip_thinking, to_pil
+from .base import Backend, GenRequest, GenResult, ScoreRequest
+
+
+class _RowStop:
+    """Per-row stopping (each request keeps its own max_new_tokens)."""
+
+    def __init__(self, prompt_len: int, limits: list[int], eos: set[int]):
+        self.prompt_len, self.limits, self.eos = prompt_len, torch.tensor(limits), eos
+
+    def __call__(self, input_ids, scores, **kw):
+        n = input_ids.shape[1] - self.prompt_len
+        done = self.limits.to(input_ids.device) <= n
+        if self.eos:
+            last = input_ids[:, -1]
+            for e in self.eos:
+                done |= last == e
+        return done
+
+
+class HFBackend(Backend):
+    name = "hf"
+    native_image_generation = False
+    verify_think = False
+
+    def __init__(self, model, processor=None, *, device_map="auto", dtype=torch.bfloat16,
+                 max_batch_size: int = 16, max_batch_tokens: int = 65536,
+                 max_image_pixels: Optional[int] = 1024 * 1024, think: bool = False,
+                 attn_implementation: Optional[str] = None, trust_remote_code: bool = False, **load_kwargs):
+        import os
+        # FP8 checkpoints: use transformers' Triton kernels, not DeepGEMM (which JIT-compiles with the system nvcc).
+        os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")
+        from transformers import AutoProcessor
+        if isinstance(model, str):
+            self.model_id = model
+            kw = dict(dtype=dtype, device_map=device_map, trust_remote_code=trust_remote_code, **load_kwargs)
+            if attn_implementation:
+                kw["attn_implementation"] = attn_implementation
+            self.model = _load_model(model, kw)
+            self.processor = processor or AutoProcessor.from_pretrained(model, trust_remote_code=trust_remote_code)
+        else:
+            if processor is None:
+                raise ValueError("pass the model's processor along with an already-loaded model")
+            self.model, self.processor = model, processor
+            self.model_id = getattr(model.config, "_name_or_path", "model")
+        self.model.eval()
+        self.tokenizer = getattr(self.processor, "tokenizer", self.processor)
+        self.tokenizer.padding_side = "left"
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.think = think
+        self.max_batch_size = max_batch_size
+        self.max_batch_tokens = max_batch_tokens
+        self.max_image_pixels = max_image_pixels
+        gc = getattr(self.model, "generation_config", None)
+        eos = getattr(gc, "eos_token_id", None) if gc is not None else None
+        eos = eos if isinstance(eos, list) else [eos] if eos is not None else []
+        if self.tokenizer.eos_token_id is not None:
+            eos.append(self.tokenizer.eos_token_id)
+        self.eos_ids = sorted({e for e in eos if e is not None})
+        self.device = next(self.model.parameters()).device
+        try:
+            self._logits_to_keep = "logits_to_keep" in inspect.signature(self.model.forward).parameters
+        except (TypeError, ValueError):
+            self._logits_to_keep = False
+        self._queue: "queue.Queue" = queue.Queue()
+        self._thread = threading.Thread(target=self._loop, daemon=True, name="simit-hf")
+        self._thread.start()
+
+    # ------------------------------------------------------------- prompts
+    def _image(self, img: Image.Image) -> Image.Image:
+        img = to_pil(img)
+        if self.max_image_pixels and img.width * img.height > self.max_image_pixels:
+            s = (self.max_image_pixels / (img.width * img.height)) ** 0.5
+            img = img.resize((max(28, int(img.width * s)), max(28, int(img.height * s))), Image.BICUBIC)
+        return img
+
+    def _render(self, parts, think: bool, add_generation_prompt: bool = True):
+        content, images = [], []
+        for p in parts:
+            if isinstance(p, Image.Image):
+                content.append({"type": "image"})
+                images.append(self._image(p))
+            elif content and content[-1]["type"] == "text":
+                content[-1]["text"] += p
+            else:
+                content.append({"type": "text", "text": p})
+        messages = [{"role": "user", "content": content}]
+        try:
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt,
+                                                      enable_thinking=think)
+        except TypeError:
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt)
+        return text, images
+
+    def _inputs(self, texts, images):
+        kw = dict(text=texts, return_tensors="pt", padding=True)
+        flat = [im for ims in images for im in ims]
+        if flat:
+            kw["images"] = flat
+        return self.processor(**kw).to(self.device)
+
+    # ------------------------------------------------------------ requests
+    def submit_generate(self, req: GenRequest) -> Future:
+        fut: Future = Future()
+        self._queue.put(("gen", req, fut))
+        return fut
+
+    def submit_score(self, req: ScoreRequest) -> Future:
+        fut: Future = Future()
+        self._queue.put(("score", req, fut))
+        return fut
+
+    def close(self):
+        self._queue.put(None)
+
+    # --------------------------------------------------------------- worker
+    def _loop(self):
+        torch.set_grad_enabled(False)
+        while True:
+            item = self._queue.get()
+            if item is None:
+                return
+            items = [item]
+            try:
+                while True:
+                    nxt = self._queue.get_nowait()
+                    if nxt is None:
+                        self._queue.put(None)
+                        break
+                    items.append(nxt)
+            except queue.Empty:
+                pass
+            items = [it for it in items if not it[2].cancelled()]
+            gens = [it for it in items if it[0] == "gen"]
+            scores = [it for it in items if it[0] == "score"]
+            # Long generations (code, captions) must not hold short ones hostage.
+            groups: dict = {}
+            for it in gens:
+                r = it[1]
+                key = (r.temperature > 0, round(r.temperature, 3), r.top_p, r.repetition_penalty,
+                       r.think or self.think, r.logprobs, r.max_new_tokens > 256)
+                groups.setdefault(key, []).append(it)
+            # Cheap work first (scoring, short answers like routing), long generations last.
+            for chunk in self._chunks(scores):
+                self._safe(self._run_score, chunk)
+            for key, group in sorted(groups.items(), key=lambda kv: max(it[1].max_new_tokens for it in kv[1])):
+                for chunk in self._chunks(group):
+                    self._safe(self._run_gen, chunk)
+
+    def _cost(self, req) -> int:
+        """Rough prompt length in tokens (images at ~one token per 28x28 pixels)."""
+        n = 0
+        for p in req.parts:
+            if isinstance(p, Image.Image):
+                px = p.width * p.height
+                n += min(px, self.max_image_pixels or px) // 784 + 1
+            else:
+                n += len(p) // 3 + 1
+        return n + len(getattr(req, "target", "")) // 3 + getattr(req, "max_new_tokens", 0)
+
+    def _chunks(self, items):
+        """Batches of at most ``max_batch_size`` requests and ``max_batch_tokens`` tokens."""
+        batch, tokens = [], 0
+        for it in items:
+            c = self._cost(it[1])
+            if batch and (len(batch) >= self.max_batch_size or tokens + c > self.max_batch_tokens):
+                yield batch
+                batch, tokens = [], 0
+            batch.append(it)
+            tokens += c
+        if batch:
+            yield batch
+
+    def _safe(self, fn, items):
+        try:
+            fn(items)
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+            # cuBLAS reports a failed workspace allocation as its own error
+            if not isinstance(e, torch.cuda.OutOfMemoryError) and "CUBLAS_STATUS_" not in str(e):
+                tb = traceback.format_exc()
+                for _, _, fut in items:
+                    if not fut.done():
+                        fut.set_exception(RuntimeError(f"{e}\n{tb}"))
+                return
+            torch.cuda.empty_cache()
+            if len(items) > 1:  # retry in halves, and keep later batches below this size
+                self.max_batch_tokens = max(2048, min(self.max_batch_tokens,
+                                                      sum(self._cost(it[1]) for it in items) // 2))
+                mid = len(items) // 2
+                self._safe(fn, items[:mid])
+                self._safe(fn, items[mid:])
+            else:
+                items[0][2].set_exception(RuntimeError("CUDA out of memory for a single request"))
+        except Exception as e:
+            tb = traceback.format_exc()
+            for _, _, fut in items:
+                if not fut.done():
+                    fut.set_exception(RuntimeError(f"{e}\n{tb}"))
+
+    def _run_gen(self, items):
+        reqs = [r for _, r, _ in items]
+        rendered = [self._render(r.parts, think=r.think or self.think) for r in reqs]
+        inputs = self._inputs([t for t, _ in rendered], [ims for _, ims in rendered])
+        r0 = reqs[0]
+        prompt_len = inputs["input_ids"].shape[1]
+        limits = [r.max_new_tokens for r in reqs]
+        from transformers import StoppingCriteriaList
+        gen_kw = dict(max_new_tokens=max(limits), return_dict_in_generate=True, output_logits=r0.logprobs,
+                      stopping_criteria=StoppingCriteriaList([_RowStop(prompt_len, limits, set(self.eos_ids))]),
+                      pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.eos_ids or None)
+        if r0.temperature > 0:
+            gen_kw.update(do_sample=True, temperature=r0.temperature, top_p=r0.top_p, top_k=0)
+        else:
+            gen_kw.update(do_sample=False, temperature=None, top_p=None, top_k=None)
+        if r0.repetition_penalty not in (None, 1.0):
+            gen_kw["repetition_penalty"] = r0.repetition_penalty
+        out = self.model.generate(**inputs, **gen_kw)
+        seqs = out.sequences[:, prompt_len:]
+        lps = None
+        if r0.logprobs:
+            logits = torch.stack(out.logits, dim=1).float()            # [B, T, V]
+            lps = torch.log_softmax(logits, dim=-1).gather(-1, seqs[:, :logits.shape[1], None]).squeeze(-1)
+        for i, ((_, r, fut), row) in enumerate(zip(items, seqs)):
+            ids = row.tolist()
+            n = len(ids)
+            for j, t in enumerate(ids):
+                if t in self.eos_ids or t == self.tokenizer.pad_token_id and j > 0 and ids[j - 1] in self.eos_ids:
+                    n = j
+                    break
+            n = min(n, r.max_new_tokens)
+            ids = ids[:n]
+            raw = self.tokenizer.decode(ids, skip_special_tokens=False)
+            text = self.tokenizer.decode(ids, skip_special_tokens=True)
+            text = strip_thinking(text) if (r.think or self.think or "<think>" in text or "</think>" in text) else text.strip()
+            res = GenResult(text=text, raw_text=raw, num_tokens=n,
+                            token_logprobs=lps[i, :n].tolist() if lps is not None else None)
+            if not fut.done():
+                fut.set_result(res)
+
+    def _run_score(self, items):
+        reqs = [r for _, r, _ in items]
+        prompts, images, targets = [], [], []
+        for r in reqs:
+            text, ims = self._render(r.parts, think=False)
+            prompts.append(text)
+            images.append(ims)
+            targets.append(self.tokenizer(r.target, add_special_tokens=False)["input_ids"])
+        # Score "prompt + target" with right-aligned targets: append target ids
+        # to each (left-padded) prompt, then read logits one position earlier.
+        enc = self._inputs(prompts, images)
+        ids, mask = enc["input_ids"], enc["attention_mask"]
+        T = max(len(t) for t in targets)
+        pad = self.tokenizer.pad_token_id
+        tgt = torch.full((len(reqs), T), pad, dtype=ids.dtype, device=ids.device)
+        tmask = torch.zeros((len(reqs), T), dtype=mask.dtype, device=mask.device)
+        for i, t in enumerate(targets):
+            if t:
+                tgt[i, :len(t)] = torch.tensor(t, device=ids.device)
+                tmask[i, :len(t)] = 1
+        for key, val in list(enc.items()):  # other per-token fields (token types...) extend with zeros
+            if key not in ("input_ids", "attention_mask") and torch.is_tensor(val) and val.shape == ids.shape:
+                enc[key] = torch.cat([val, torch.zeros_like(tgt, dtype=val.dtype)], dim=1)
+        enc["input_ids"] = torch.cat([ids, tgt], dim=1)
+        enc["attention_mask"] = torch.cat([mask, tmask], dim=1)
+        P = ids.shape[1]
+        if self._logits_to_keep:  # only the T+1 positions that predict target tokens, not the whole prompt
+            logits = self.model(**enc, logits_to_keep=T + 1).logits
+            start = 0
+        else:
+            logits = self.model(**enc).logits
+            start = P - 1
+        for i, ((_, r, fut), t) in enumerate(zip(items, targets)):
+            if not t:
+                fut.set_result([])
+                continue
+            lg = logits[i, start:start + len(t)].float()
+            lp = torch.log_softmax(lg, dim=-1).gather(-1, torch.tensor(t, device=lg.device)[:, None]).squeeze(-1)
+            fut.set_result(lp.tolist())
+
+
+def _dense_fp8_fix(model_id: str, kw: dict) -> None:
+    """Pre-quantized FP8 checkpoints of dense models can list MoE router names
+    ("...mlp.gate") as modules not to convert; transformers matches those as
+    prefixes, which also skips the dense "...mlp.gate_proj" and loads its FP8
+    weights unscaled (garbage outputs). Drop them when the model has no experts."""
+    if "config" in kw:
+        return
+    try:
+        from transformers import AutoConfig
+        config = AutoConfig.from_pretrained(model_id, trust_remote_code=kw.get("trust_remote_code", False))
+    except Exception:
+        return
+    quant = getattr(config, "quantization_config", None)
+    text = getattr(config, "text_config", config)
+    if not isinstance(quant, dict) or not quant.get("modules_to_not_convert") or getattr(text, "num_experts", 0):
+        return
+    keep = [m for m in quant["modules_to_not_convert"] if not m.endswith(".mlp.gate")]
+    if len(keep) != len(quant["modules_to_not_convert"]):
+        quant["modules_to_not_convert"] = keep
+        kw["config"] = config
+
+
+def _load_model(model_id: str, kw: dict):
+    import transformers
+    _dense_fp8_fix(model_id, kw)
+    errors = []
+    for name in ("AutoModelForMultimodalLM", "AutoModelForImageTextToText", "AutoModelForVision2Seq",
+                 "AutoModelForCausalLM"):
+        cls = getattr(transformers, name, None)
+        if cls is None:
+            continue
+        try:
+            return cls.from_pretrained(model_id, **kw)
+        except (ValueError, KeyError) as e:  # architecture not registered for this auto class
+            errors.append(f"{name}: {str(e).splitlines()[0]}")
+    raise ValueError(f"could not load {model_id} with transformers: " + "; ".join(errors))

simit/backends/imagegen.py

@@ -0,0 +1,136 @@
+"""Text-to-image tools for models that cannot generate images themselves."""
+
+from __future__ import annotations
+
+import inspect
+import queue
+import threading
+import traceback
+from concurrent.futures import Future
+from typing import Callable, Optional
+
+import torch
+from PIL import Image
+
+from .base import Backend, ImageRequest
+
+# Sampler settings for distilled checkpoints the defaults would waste steps on.
+_PRESETS = {
+    "flux.2-klein": dict(num_inference_steps=4, guidance_scale=1.0),
+    "flux.1-schnell": dict(num_inference_steps=4, guidance_scale=0.0),
+    "sdxl-turbo": dict(num_inference_steps=2, guidance_scale=0.0),
+    "sd-turbo": dict(num_inference_steps=2, guidance_scale=0.0),
+}
+
+
+class DiffusersImageGenerator(Backend):
+    """Wraps a diffusers text-to-image pipeline; requests of equal size are
+    generated in one batch."""
+
+    name = "diffusers"
+    native_image_generation = True
+
+    def __init__(self, model, device: Optional[str] = None, dtype=torch.bfloat16, max_batch_size: int = 8,
+                 **call_kwargs):
+        if isinstance(model, str):
+            from diffusers import DiffusionPipeline
+            from diffusers.pipelines import pipeline_utils
+            self.model_id = model
+            dtype_kw = "dtype" if hasattr(pipeline_utils, "_resolve_dtype") else "torch_dtype"  # renamed in diffusers 0.40
+            self.pipe = DiffusionPipeline.from_pretrained(model, **{dtype_kw: dtype})
+            self.pipe.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        else:
+            self.pipe = model
+            self.model_id = getattr(getattr(model, "config", None), "_name_or_path", type(model).__name__)
+        self.pipe.set_progress_bar_config(disable=True)
+        preset = next((v for k, v in _PRESETS.items() if k in self.model_id.lower()), {})
+        self.call_kwargs = {**preset, **call_kwargs}
+        self.max_batch_size = max_batch_size
+        self._queue: "queue.Queue" = queue.Queue()
+        self._thread = threading.Thread(target=self._loop, daemon=True, name="simit-imagegen")
+        self._thread.start()
+
+    def submit_image(self, req: ImageRequest) -> Future:
+        fut: Future = Future()
+        self._queue.put((req, fut))
+        return fut
+
+    def close(self):
+        self._queue.put(None)
+
+    def _loop(self):
+        torch.set_grad_enabled(False)
+        while True:
+            item = self._queue.get()
+            if item is None:
+                return
+            items = [item]
+            try:
+                while len(items) < 64:
+                    nxt = self._queue.get_nowait()
+                    if nxt is None:
+                        self._queue.put(None)
+                        break
+                    items.append(nxt)
+            except queue.Empty:
+                pass
+            by_size: dict = {}
+            for req, fut in items:
+                if not fut.cancelled():
+                    by_size.setdefault((_round(req.width), _round(req.height)), []).append((req, fut))
+            for (w, h), group in by_size.items():
+                for i in range(0, len(group), self.max_batch_size):
+                    self._run(group[i:i + self.max_batch_size], w, h)
+
+    def _run(self, group, w, h):
+        try:
+            gens = None
+            if any(r.seed is not None for r, _ in group):
+                dev = getattr(self.pipe, "device", "cpu")
+                gens = [torch.Generator(device=dev).manual_seed(r.seed if r.seed is not None else torch.seed() % 2 ** 31)
+                        for r, _ in group]
+            images = self.pipe(prompt=[r.prompt for r, _ in group], width=w, height=h, generator=gens,
+                               **self.call_kwargs).images
+            for (_, fut), im in zip(group, images):
+                if not fut.done():
+                    fut.set_result(im.convert("RGB"))
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            if len(group) > 1:  # e.g. a VLM engine sharing the GPU: retry in halves
+                mid = len(group) // 2
+                self._run(group[:mid], w, h)
+                self._run(group[mid:], w, h)
+            else:
+                for _, fut in group:
+                    if not fut.done():
+                        fut.set_exception(RuntimeError("CUDA out of memory generating a single image"))
+        except Exception as e:
+            tb = traceback.format_exc()
+            for _, fut in group:
+                if not fut.done():
+                    fut.set_exception(RuntimeError(f"{e}\n{tb}"))
+
+
+class CallableImageGenerator(Backend):
+    """Any ``fn(prompt, width, height) -> PIL.Image`` (e.g. an API client)."""
+
+    name = "callable"
+    native_image_generation = True
+
+    def __init__(self, fn: Callable):
+        self.fn = fn
+        self._pool = None
+
+    def submit_image(self, req: ImageRequest) -> Future:
+        from concurrent.futures import ThreadPoolExecutor
+        if self._pool is None:
+            self._pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix="simit-imagefn")
+        return self._pool.submit(self.fn, req.prompt, req.width, req.height)
+
+    def close(self):
+        if self._pool is not None:
+            self._pool.shutdown(wait=False)
+
+
+def _round(v: int, multiple: int = 16) -> int:
+    return max(multiple, int(round(v / multiple)) * multiple)