simit 0.1.0__py3-none-any.whl

simit/backends/lance.py

@@ -0,0 +1,338 @@
+"""Lance-3B (bytedance-research/Lance) on the batched MoT engine.
+
+Lance pairs BAGEL's Mixture-of-Transformers trunk with the Qwen2.5-VL vision
+tower, multimodal (t, h, w) RoPE and the Wan2.2 VAE. Prompt formats and
+decoding follow the setup the SIMIT paper evaluated Lance with: chat-template
+turns with task system prompts, a newline start token inside the assistant
+turn, sampling (T >= 0.8, repetition penalty 1.05) for free generation and
+greedy decoding for confidence scores.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+from PIL import Image
+
+from ..utils import image_hash, strip_thinking, to_pil
+from .base import Backend, GenRequest, GenResult, ImageRequest, ScoreRequest, chained
+from .mot.engine import ImageChunk, ImageJob, MoTEngine, TextChunk, TextJob, auto_kv_tokens
+from .mot.modeling import DeviceMixin, MoTLanguageModel, MRotaryEmbedding, PositionTable, TimestepEmbedder
+
+X2T_SYSTEM_PROMPT = "Look at the image carefully and answer the question."
+TEXT_SYSTEM_PROMPT = "You are a helpful assistant."
+T2I_SYSTEM_PROMPT = ("Describe the image by detailing the color, quantity, text, shape, size, texture, "
+                     "spatial relationships of the objects and background:")
+
+_VIT_BUCKET_RESOLUTION = 672
+_VIT_BUCKET_STRIDE = 16
+_VIT_DIVISIBLE_CROP = 28
+_VIT_ASPECT_RATIOS = ("21:9", "16:9", "4:3", "1:1", "3:4", "9:16")
+_VIT_T_SHIFT = 1000   # ViT tokens live far away on the t axis of mRoPE
+
+
+def _vit_buckets():
+    import math
+    max_area = _VIT_BUCKET_RESOLUTION ** 2
+    out = []
+    for name in _VIT_ASPECT_RATIOS:
+        ws, hs = (int(v) for v in name.split(":"))
+        aspect = ws / hs
+        bw1 = round(math.sqrt(max_area * aspect) / _VIT_BUCKET_STRIDE) * _VIT_BUCKET_STRIDE
+        bh1 = round(bw1 / aspect / _VIT_BUCKET_STRIDE) * _VIT_BUCKET_STRIDE
+        bh2 = round(math.sqrt(max_area / aspect) / _VIT_BUCKET_STRIDE) * _VIT_BUCKET_STRIDE
+        bw2 = round(bh2 * aspect / _VIT_BUCKET_STRIDE) * _VIT_BUCKET_STRIDE
+        d1, d2 = abs(bw1 / bh1 - aspect), abs(bw2 / bh2 - aspect)
+        if d1 < d2 or (d1 == d2 and abs(bh1 * bw1 - max_area) <= abs(bh2 * bw2 - max_area)):
+            out.append(((bh1, bw1), bw1 / bh1))
+        else:
+            out.append(((bh2, bw2), bw2 / bh2))
+    return out
+
+
+_BUCKETS = _vit_buckets()
+
+
+def _bucket_resize(img: Image.Image) -> Image.Image:
+    """Lance's ViT preprocessing: nearest aspect bucket (center crop + bicubic
+    resize to ~672^2 px), then center-crop to multiples of 28."""
+    from torchvision.transforms import InterpolationMode, RandomResizedCrop
+    ratio = img.width / img.height
+    (bh, bw), br = min(_BUCKETS, key=lambda b: abs(ratio - b[1]))
+    out = RandomResizedCrop(size=(bh, bw), scale=(1.0, 1.0), ratio=(br, br),
+                            interpolation=InterpolationMode.BICUBIC)(img)
+    w, h = out.size
+    cw, ch = w - w % _VIT_DIVISIBLE_CROP, h - h % _VIT_DIVISIBLE_CROP
+    if (cw, ch) != (w, h):
+        left, top = (w - cw) // 2, (h - ch) // 2
+        out = out.crop((left, top, left + cw, top + ch))
+    return out
+
+
+class LanceModel(DeviceMixin):
+    latent_channels = 48
+    latent_downsample = 16
+    max_latent_side = 64
+    num_timesteps = 30
+    timestep_shift = 3.5
+    cfg_text_scale = 4.0
+    cfg_interval = (0.4, 1.0)
+    cfg_renorm_min = 0.0
+
+    @classmethod
+    def load(cls, model_id: str = "bytedance-research/Lance", device: str = "cuda",
+             dtype=torch.bfloat16) -> "LanceModel":
+        """From a Hub id (downloaded on first use) or a local path."""
+        path = Path(model_id)
+        if not (path / "Lance_3B").exists():
+            from huggingface_hub import snapshot_download
+            path = Path(snapshot_download(model_id, allow_patterns=[
+                "config.json", "Lance_3B/*", "Qwen2.5-VL-ViT/*", "Wan2.2_VAE.pth"]))
+        model = cls(str(path), device=device, dtype=dtype)
+        model.model_id = model_id
+        return model
+
+    def __init__(self, root: str, device: str = "cuda", dtype=torch.bfloat16):
+        self.model_id = str(root)
+        from safetensors import safe_open
+        from transformers import AutoTokenizer, Qwen2VLImageProcessor
+        from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
+        from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel
+
+        from .mot.wan_vae import LanceWanVAE
+
+        self.device, self.dtype = torch.device(device), dtype
+        root = Path(root)
+        ckpt = root / "Lance_3B"
+        cfg = json.loads((ckpt / "llm_config.json").read_text())
+        hidden = cfg["hidden_size"]
+        self.latent_dim = self.latent_channels
+        with torch.device("meta"):
+            self.lm = MoTLanguageModel(cfg)
+            self.latent_pos_embed = PositionTable(self.max_latent_side, hidden)
+            self.time_embedder = TimestepEmbedder(hidden)
+            self.vae2llm = torch.nn.Linear(self.latent_dim, hidden)
+            self.llm2vae = torch.nn.Linear(hidden, self.latent_dim)
+        with safe_open(str(ckpt / "model.safetensors"), framework="pt", device=str(self.device)) as f:
+            keys = set(f.keys())
+            for prefix, mod in (("language_model", self.lm), ("latent_pos_embed", self.latent_pos_embed),
+                                ("time_embedder", self.time_embedder), ("vae2llm", self.vae2llm),
+                                ("llm2vae", self.llm2vae)):
+                sd = {n: f.get_tensor(f"{prefix}.{n}").to(dtype) for n, _ in mod.named_parameters()
+                      if f"{prefix}.{n}" in keys}
+                if "lm_head.weight" not in sd and prefix == "language_model":  # tied embeddings
+                    sd["lm_head.weight"] = sd["model.embed_tokens.weight"]
+                missing = [n for n, _ in mod.named_parameters() if n not in sd]
+                if missing:
+                    raise RuntimeError(f"Lance checkpoint is missing {prefix} weights: {missing[:5]}")
+                mod.load_state_dict(sd, strict=False, assign=True)
+                mod.eval()
+        head_dim = hidden // cfg["num_attention_heads"]
+        sections = cfg.get("rope_scaling", {}).get("mrope_section", [16, 24, 24])
+        self.lm.rotary = MRotaryEmbedding(head_dim, cfg.get("rope_theta", 1e6), sections).to(self.device)
+
+        vit_cfg = json.loads((root / "Qwen2.5-VL-ViT" / "config.json").read_text())
+        vit_cfg["_attn_implementation"] = "sdpa"
+        self.vit = Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLVisionConfig(**vit_cfg))
+        with safe_open(str(root / "Qwen2.5-VL-ViT" / "vit.safetensors"), framework="pt", device="cpu") as f:
+            sd = {k: f.get_tensor(k) for k in f.keys()}
+        missing, unexpected = self.vit.load_state_dict(sd, strict=False)
+        if [m for m in missing if "rotary" not in m]:
+            raise RuntimeError(f"Lance ViT checkpoint mismatch: missing {missing[:5]}")
+        self.vit.to(self.device, dtype).eval()
+        self.merge = vit_cfg.get("spatial_merge_size", 2)
+        self.image_processor = Qwen2VLImageProcessor()
+
+        self.vae = LanceWanVAE(str(root / "Wan2.2_VAE.pth"), dtype=torch.bfloat16, device=self.device)
+
+        tok = AutoTokenizer.from_pretrained(str(ckpt))
+        self.tokenizer = tok
+        self.bos = tok.convert_tokens_to_ids("<|im_start|>")
+        self.eos = tok.convert_tokens_to_ids("<|im_end|>")
+        self.start_of_image = tok.convert_tokens_to_ids("<|vision_start|>")
+        self.end_of_image = tok.convert_tokens_to_ids("<|vision_end|>")
+        self.stop_ids = (self.eos, tok.convert_tokens_to_ids("<|endoftext|>"))
+        nl = tok.encode("\n", add_special_tokens=False)
+        assert len(nl) == 1
+        self.newline = nl[0]
+
+    # ---------------------------------------------------------- understanding
+    def preprocess_image(self, image) -> ImageChunk:
+        img = _bucket_resize(to_pil(image))
+        out = self.image_processor(images=img, return_tensors="pt")
+        pix, grid = out["pixel_values"], out["image_grid_thw"]
+        t, h, w = (int(v) for v in grid[0].tolist())
+        n = t * h * w // (self.merge ** 2)
+        return ImageChunk(n=n + 2, key=image_hash(img), payload=(pix, grid, h // self.merge, w // self.merge))
+
+    def encode_images(self, payloads):
+        pix = torch.cat([p[0] for p in payloads]).to(self.device, self.dtype)
+        grid = torch.cat([p[1] for p in payloads]).to(self.device)
+        out = self.vit(hidden_states=pix, grid_thw=grid)
+        if torch.is_tensor(out):
+            merged = out
+        else:  # BaseModelOutputWithPooling: the merger output is the pooled one
+            merged = getattr(out, "pooler_output", None)
+            if merged is None:
+                merged = out[0]
+        sizes = [p[2] * p[3] for p in payloads]
+        return list(merged.to(self.dtype).split(sizes))
+
+    # (t, h, w) rope positions
+    def text_positions(self, start: int, n: int):
+        return [(p, p, p) for p in range(start, start + n)]
+
+    def image_positions(self, start: int, chunk):
+        _, _, hm, wm = chunk.payload
+        m = max(hm, wm)
+        pos = [(_VIT_T_SHIFT, start, start)]
+        pos += [(_VIT_T_SHIFT + 1, start + 1 + i, start + 1 + j) for i in range(hm) for j in range(wm)]
+        pos.append((_VIT_T_SHIFT + m + 1, start + m + 1, start + m + 1))
+        return pos
+
+    def position_after_image(self, pos: int, chunk) -> int:
+        _, _, hm, wm = chunk.payload
+        return pos + max(hm, wm) + 2
+
+    def position_after(self, chunks, n: int) -> int:
+        return n  # the prefix cache only covers leading text
+
+    # -------------------------------------------------------------- generation
+    def latent_positions(self, start: int, h: int, w: int):
+        m = max(h, w)
+        pos = [(start, start, start)]
+        pos += [(start + 1, start + 1 + i, start + 1 + j) for i in range(h) for j in range(w)]
+        pos.append((start + m + 1,) * 3)
+        return pos
+
+    def latent_pos_ids(self, h: int, w: int):
+        return (torch.arange(h)[:, None] * self.max_latent_side + torch.arange(w)).flatten()
+
+    def denoise_schedule(self):
+        n, s = self.num_timesteps, self.timestep_shift
+        t = torch.linspace(1, 0, n + 1)  # Lance runs exactly `num_timesteps` Euler steps
+        t = s * t / (1 + (s - 1) * t)
+        return [(float(t[i]), float(t[i] - t[i + 1])) for i in range(n)]
+
+    def cfg_active(self, t: float) -> bool:
+        lo, hi = self.cfg_interval
+        return lo < t <= hi and self.cfg_text_scale > 1.0
+
+    def embed_latents(self, x, t: float, pos_ids):
+        tt = torch.full((x.shape[0],), t, device=x.device)
+        out = self.vae2llm(x.to(self.dtype)) + self.time_embedder(tt) + self.latent_pos_embed(pos_ids).to(self.dtype)
+        return out.to(self.dtype)
+
+    def latents_to_velocity(self, h):
+        return self.llm2vae(h)
+
+    def apply_cfg(self, v, v_uncond):
+        v, v_uncond = v.float(), v_uncond.float()
+        guided = v_uncond + self.cfg_text_scale * (v - v_uncond)
+        scale = (v.norm() / (guided.norm() + 1e-8)).clamp(min=self.cfg_renorm_min, max=1.0)
+        return guided * scale
+
+    def decode_latents(self, items):
+        out = []
+        for x, h, w in items:
+            lat = x.reshape(1, h, w, self.latent_channels).permute(0, 3, 1, 2)
+            img = self.vae.decode(lat.to(torch.bfloat16)).float()
+            img = ((img * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
+            out.append(Image.fromarray(img))
+        return out
+
+
+class LanceBackend(Backend):
+    name = "lance"
+    native_image_generation = True
+    verify_think = False
+    #: Lance follows short single-turn instructions only (see the paper's
+    #: cross-model study): triplets are synthesized in three simple calls,
+    #: every description is realized by native generation, no critic.
+    pipeline_defaults = {"synthesis": "decomposed", "use_skills": False, "verify": False}
+    #: Greedy decoding degenerates on this checkpoint; free generation samples.
+    min_temperature = 0.8
+    repetition_penalty = 1.05
+
+    def __init__(self, model_id: Union[str, "LanceModel"] = "bytedance-research/Lance", device: Optional[str] = None,
+                 kv_cache_tokens: Optional[int] = None, max_step_tokens: int = 16384,
+                 image_steps: Optional[int] = None, use_cuda_graphs: bool = True):
+        """``model_id``: a Hub id / local path, or a :class:`LanceModel` loaded earlier
+        (e.g. on CPU at startup, then moved per request). ``device``: where to put
+        the weights (default ``cuda:0``; for a loaded model, where it already is)."""
+        if isinstance(model_id, LanceModel):
+            self.model = model_id
+            if device is not None and self.model.device != torch.device(device):
+                self.model.to(device)
+        else:
+            self.model = LanceModel.load(model_id, device=device or "cuda:0")
+        self.model_id = self.model.model_id
+        if image_steps is not None:  # fewer denoising steps: faster, slightly lower fidelity
+            self.model.num_timesteps = int(image_steps)
+        if kv_cache_tokens is None:
+            kv_cache_tokens = auto_kv_tokens(self.model)
+        self.engine = MoTEngine(self.model, kv_tokens=kv_cache_tokens, max_step_tokens=max_step_tokens,
+                                use_cuda_graphs=use_cuda_graphs)
+
+    def _encode(self, text: str):
+        return self.model.tokenizer.encode(text, add_special_tokens=False)
+
+    def _chunks(self, parts):
+        m = self.model
+        system = X2T_SYSTEM_PROMPT if any(isinstance(p, Image.Image) for p in parts) else TEXT_SYSTEM_PROMPT
+        chunks, buf = [], f"<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n"
+        for p in parts:
+            if isinstance(p, Image.Image):
+                chunks.append(TextChunk(self._encode(buf)))
+                chunks.append(m.preprocess_image(p))
+                buf = ""
+            else:
+                buf += p
+        buf += "<|im_end|>\n<|im_start|>assistant"
+        chunks.append(TextChunk(self._encode(buf) + [m.newline]))  # start inside the assistant turn
+        return [c for c in chunks if not isinstance(c, TextChunk) or c.ids]
+
+    def submit_generate(self, req: GenRequest):
+        if req.logprobs:  # confidence: greedy, as the paper's Lance scores were computed
+            temperature, rp = 0.0, 1.0
+        else:
+            temperature = max(req.temperature, self.min_temperature)
+            rp = req.repetition_penalty or self.repetition_penalty
+        job = TextJob(chunks=self._chunks(req.parts), max_new_tokens=req.max_new_tokens, temperature=temperature,
+                      top_p=req.top_p, repetition_penalty=rp, logprobs=req.logprobs, stop_ids=self.model.stop_ids)
+        tok = self.model.tokenizer
+        if req.on_text is not None:
+            job.on_tokens = lambda ids: req.on_text(tok.decode(ids, skip_special_tokens=True))
+
+        def post(seq):
+            raw = tok.decode(seq.out_ids, skip_special_tokens=True)
+            text = strip_thinking(raw) if req.think else raw.strip()
+            return GenResult(text=text, raw_text=raw, token_logprobs=seq.out_lp if req.logprobs else None,
+                             num_tokens=len(seq.out_ids))
+
+        return chained(self.engine.submit(job), post)
+
+    def submit_score(self, req: ScoreRequest):
+        target = self._encode(req.target)
+        if not target:
+            from concurrent.futures import Future
+            f = Future()
+            f.set_result([])
+            return f
+        chunks = self._chunks(req.parts) + [TextChunk(target)]
+        job = TextJob(chunks=chunks, score_len=len(target), stop_ids=self.model.stop_ids)
+        return chained(self.engine.submit(job), lambda seq: seq.score_lp)
+
+    def submit_image(self, req: ImageRequest):
+        m = self.model
+        w, h = (max(16, int(round(v / 16)) * 16) for v in (req.width, req.height))
+        text = (f"system\n{T2I_SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{req.prompt}<|im_end|>\n"
+                f"<|im_start|>assistant\n")
+        job = ImageJob(chunks=[TextChunk([m.bos] + self._encode(text) + [m.eos])], height=h, width=w, seed=req.seed)
+        return self.engine.submit(job)
+
+    def close(self):
+        self.engine.close()