publish: preserve nested overlay dirs (patches/ still flattened; --gist refused for nested); dsv4-vision: full OpenAI MM processor trio (info/processor/dummy+updates), model embed_multimodal + embed_input_ids merge, placeholder id 129264; verified sentinel/perm math in-image

docker-images/dsv4-flash-vision/plugin/__init__.py

@@ -1,16 +0,0 @@
-"""vLLM general plugin for DeepSeek-V4-Flash-Vision-Exp."""
-
-__all__ = ["register"]
-
-
-def register() -> None:
-    from vllm.model_executor.models.registry import ModelRegistry
-
-    ModelRegistry.register_model(
-        "DeepseekV4ForCausalLM",
-        "dsv4_vision.model:DeepseekV4VisionForCausalLM",
-    )
-    ModelRegistry.register_model(
-        "DeepseekV4VisionForCausalLM",
-        "dsv4_vision.model:DeepseekV4VisionForCausalLM",
-    )

docker-images/dsv4-flash-vision/plugin/dsv4_vision/model.py

@@ -12,9 +12,18 @@ from typing import Any
 import torch
 import torch.nn as nn
 
+import torch
+import torch.nn as nn
+
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.models.deepseek_v4 import DeepseekV4ForCausalLM
+from vllm.multimodal import MULTIMODAL_REGISTRY
 
+from .processor import (
+    DSV4VisionDummyInputsBuilder,
+    DSV4VisionMultiModalProcessor,
+    DSV4VisionProcessingInfo,
+)
 from .vision import Aligner, ViT
 
 
@@ -35,12 +44,27 @@ def _vision_args(hf: Any) -> SimpleNamespace:
     )
 
 
+@MULTIMODAL_REGISTRY.register_processor(
+    DSV4VisionMultiModalProcessor,
+    info=DSV4VisionProcessingInfo,
+    dummy_inputs=DSV4VisionDummyInputsBuilder,
+)
 class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
     # ModelConfig.is_multimodal_model keys off this class attribute
     # (interfaces.supports_multimodal) — without it the OpenAI layer
     # answers "not a multimodal model" for image_url requests.
     supports_multimodal = True
 
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        # chat_utils asks the resolved model class for the per-image prompt
+        # placeholder; the tokenizer encodes it to <｜deepseek_image｜> (id 129264).
+        if modality.startswith("image"):
+            from .encoding_dsv4 import IMAGE_PLACEHOLDER
+
+            return IMAGE_PLACEHOLDER
+        return None
+
     def __init__(self, *, vllm_config, prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         hf = vllm_config.model_config.hf_config
@@ -60,6 +84,90 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
     def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
         return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)
 
+    def _merge_mm_embeddings(
+        self,
+        inputs_embeds: torch.Tensor,
+        multimodal_embeddings,
+        is_multimodal: torch.Tensor,
+    ) -> torch.Tensor:
+        # Mirrors interfaces._merge_multimodal_embeddings: fill masked rows in
+        # order with the concatenated per-item embedding runs.
+        flat = torch.cat([emb.reshape(-1) for emb in multimodal_embeddings])
+        inputs_embeds.masked_scatter_(
+            is_multimodal.to(inputs_embeds.device).unsqueeze(-1), flat
+        )
+        return inputs_embeds
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if is_multimodal is not None:
+            # Sentinel block ids may exceed vocab; zero them before lookup.
+            input_ids = input_ids.masked_fill(
+                is_multimodal.to(input_ids.device).unsqueeze(-1), 0
+            )
+        inputs_embeds = super().embed_input_ids(input_ids)
+        if not multimodal_embeddings:
+            return inputs_embeds
+        return self._merge_mm_embeddings(
+            inputs_embeds, multimodal_embeddings, is_multimodal
+        )
+
+    def embed_multimodal(self, **kwargs) -> tuple:
+        patches = kwargs.get("image_patches")
+        if patches is None or self.vision is None:
+            return ()
+        grids = kwargs["image_grids"].data
+        types_all = kwargs["image_types"].data
+        perm_all = kwargs["image_perm"].data
+        t_counts = kwargs["image_type_counts"].data.tolist() if hasattr(kwargs["image_type_counts"], "tolist") else kwargs["image_type_counts"].data.tolist()
+        perm_counts = kwargs["image_perm_counts"].data.tolist()
+        n_items = int(grids.shape[0])
+
+        out: list[torch.Tensor] = []
+        p_off = t_off = q_off = 0
+        start = self.image_start
+        pad = self.image_pad
+        newline = self.image_newline
+        end = self.image_end
+        for i in range(n_items):
+            n_vit_h, n_vit_w, n_llm_h, n_llm_w = grids[i].tolist()
+            p_count = n_vit_h * n_vit_w
+            t_count = int(t_counts[i])
+            q_count = int(perm_counts[i])
+            item_types = types_all[t_off : t_off + t_count].tolist()
+            item_perm = perm_all[q_off : q_off + q_count].to(torch.long)
+            if q_count != sum(t == 2 for t in item_types):
+                raise ValueError(
+                    "DSV4 vision perm/sentinel layout mismatch "
+                    f"({q_count} != {sum(t == 2 for t in item_types)})"
+                )
+            t_off += t_count
+            q_off += q_count
+            vit = self.encode_image(patches[p_off : p_off + p_count], n_vit_h, n_vit_w)
+            p_off += p_count
+            rows: list[torch.Tensor] = []
+            vi = 0
+            for t in item_types:
+                if t == 0:  # IMAGE_START
+                    rows.append(start)
+                elif t == 1:  # IMAGE_PAD
+                    rows.append(pad)
+                elif t == 2:  # IMAGE
+                    rows.append(vit[int(item_perm[vi])])
+                    vi += 1
+                elif t == 3:  # IMAGE_NEW_LINE
+                    rows.append(newline)
+                else:  # IMAGE_END
+                    rows.append(end)
+            out.append(torch.stack(rows))
+
+        return tuple(out)
+
     def forward(
         self,
         input_ids,
@@ -68,8 +176,9 @@ class DeepseekV4VisionForCausalLM(DeepseekV4ForCausalLM):
         inputs_embeds=None,
         **kwargs,
     ):
-        # Spec-decode extra kwargs must not TypeError. lm_head and
-        # get_mtp_target_hidden_states stay on the parent class.
+        # The runner embeds via embed_input_ids (mm merge) and hands us
+        # merged inputs_embeds; lm_head and get_mtp_target_hidden_states
+        # stay on the parent class.
         return super().forward(
             input_ids,
             positions,

docker-images/dsv4-flash-vision/plugin/dsv4_vision/processor.py

@@ -0,0 +1,352 @@
+"""OpenAI-multimodal processor wiring for DeepSeek-V4-Flash-Vision-Exp.
+
+dots3_note-style processor trio (workflow of the b12x vLLM fork):
+
+- a small HF-like processor that tokenizes prompt text (placeholders intact)
+  and precomputes ViT patches / grid / sentinel layout per image,
+- a dummy builder for MM profiling,
+- a `BaseMultiModalProcessor` whose prompt update expands every
+  `<｜deepseek_image|>` token (vocab id 129264) into the per-image sentinel
+  block built by `image_processor.build_image_block`, with the whole block
+  flagged as mm embeddings.
+
+`model.embed_multimodal` turns the tensors into per-image embedding runs and
+`model.embed_input_ids` scatters them over the flagged range on prefill.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from functools import cached_property
+from typing import Any, cast
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import BatchFeature
+
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdateDetails,
+)
+
+from .encoding_dsv4 import IMAGE_PLACEHOLDER
+from .image_processor import (
+    COMPRESS_PAD_TO,
+    build_image_block,
+    grid_tokens,
+    safe_resize,
+    solve_resize_ratio,
+)
+
+_PROTO_PATCHES = 60  # prototype patch-grid side for the dummy-size sweep
+
+
+def _vision_kv(hf: Any) -> dict[str, int | float]:
+    return {
+        "vision_patch_size": int(getattr(hf, "vision_patch_size", 14)),
+        "vision_downsample_ratio": int(getattr(hf, "vision_downsample_ratio", 3)),
+        "vision_max_n_token": int(getattr(hf, "vision_max_n_token", 384)),
+        "vision_min_pixels": int(getattr(hf, "vision_min_pixels", 147456)),
+        "vision_max_wh_ratio": int(getattr(hf, "vision_max_wh_ratio", 8)),
+    }
+
+
+class DSV4VisionImagePreparer:
+    """Resize/pad/patch a PIL image per the day-0 `load_image` math."""
+
+    def __init__(self, hf: Any) -> None:
+        for key, value in _vision_kv(hf).items():
+            setattr(self, key, value)
+        self.image_token_id: int | None = None
+
+    @property
+    def p(self) -> int:
+        return self.vision_patch_size
+
+    def prepare(self, image: Image.Image) -> dict[str, torch.Tensor]:
+        image = image.convert("RGB")
+        width, height = image.size
+        if self.vision_max_wh_ratio and width > height * self.vision_max_wh_ratio:
+            width = height * self.vision_max_wh_ratio
+        if 0 < width * height < self.vision_min_pixels:
+            scale = (self.vision_min_pixels / (width * height)) ** 0.5
+            width = int(width * scale)
+            height = int(height * scale)
+        best_width = math.ceil(width / self.p) * self.p
+        best_height = math.ceil(height / self.p) * self.p
+        n_llm_h, n_llm_w, best_height, best_width = safe_resize(
+            height,
+            width,
+            best_height,
+            best_width,
+            self.p,
+            self.vision_downsample_ratio,
+            self.vision_max_n_token,
+        )
+        n_vit_h, n_vit_w = best_height // self.p, best_width // self.p
+        if self.vision_max_wh_ratio and image.width >= self.vision_max_wh_ratio * image.height:
+            resized = image.resize((best_width, best_height))
+        else:
+            resized = _pad_image(image, best_width, best_height)
+        x = torch.from_numpy(np.asarray(resized, dtype="float32")).permute(2, 0, 1) / 255
+        x = ((x - 0.5) / 0.5).to(torch.bfloat16)
+        patches = (
+            x.reshape(3, n_vit_h, self.p, n_vit_w, self.p)
+            .permute(1, 3, 0, 2, 4)
+            .reshape(n_vit_h * n_vit_w, 3, self.p, self.p)
+        )
+        return dict(
+            patches=patches,
+            n_vit_h=n_vit_h,
+            n_vit_w=n_vit_w,
+            n_llm_h=n_llm_h,
+            n_llm_w=n_llm_w,
+        )
+
+    def max_block_scan(self) -> tuple[int, int, int]:
+        """(best_h, best_w) and its block length over candidate aspect ratios."""
+        base = _PROTO_PATCHES * self.p * self.vision_downsample_ratio
+        ratios = [1 / r for r in (1, 2, 3, 4, 5, 6, 7, 8)]
+        ratios += [1 / r for r in ratios[1:]]
+        best: tuple[int, int, int] = (0, 0, 0)
+        for height_ratio in ratios:
+            height = max(self.p * self.vision_downsample_ratio, int(base * height_ratio))
+            width = max(self.p * self.vision_downsample_ratio, int(base / height_ratio))
+            best_h, best_w = math.ceil(height / self.p) * self.p, math.ceil(width / self.p) * self.p
+            n_h, n_w, fb_h, fb_w, _ = solve_resize_ratio(
+                height, width, self.p, self.vision_downsample_ratio, self.vision_max_n_token
+            )
+            best_h, best_w = max(best_h, fb_h), max(best_w, fb_w)
+            _, _, num = grid_tokens(fb_h, fb_w, self.p, self.vision_downsample_ratio)
+            if num > best[2]:
+                best = (fb_h, fb_w, num + COMPRESS_PAD_TO - 1)
+        return best
+
+
+def _pad_image(image: Image.Image, width: int, height: int) -> Image.Image:
+    from PIL import ImageOps
+
+    return ImageOps.pad(image, (width, height), color=(127, 127, 127))
+
+
+class DSV4VisionProcessingInfo(BaseProcessingInfo):
+    @cached_property
+    def preparer(self) -> DSV4VisionImagePreparer:
+        return DSV4VisionImagePreparer(self.get_hf_config())
+
+    @cached_property
+    def processor(self) -> "DSV4VisionProcessor":
+        return DSV4VisionProcessor(
+            self.get_tokenizer(),
+            self.preparer,
+            vocab_size=int(self.get_hf_config().vocab_size),
+        )
+
+    def get_hf_processor(self, **kwargs: object) -> "DSV4VisionProcessor":
+        return self.processor
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"image": None}
+
+    @cached_property
+    def max_image_size(self) -> tuple[int, int]:
+        best_h, best_w, _ = self.preparer.max_block_scan()
+        return best_w, best_h
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        if not mm_counts.get("image"):
+            return {}
+        _, _, num = self.preparer.max_block_scan()
+        return {"image": num + 4}
+
+
+class DSV4VisionProcessor:
+    """HF-like processor: tokenize text, precompute image tensors."""
+
+    def __init__(self, tokenizer, preparer: DSV4VisionImagePreparer, vocab_size: int) -> None:
+        self.tokenizer = tokenizer
+        self.preparer = preparer
+        self.vocab_size = vocab_size
+        ids = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
+        if ids is None or ids == getattr(tokenizer, "unk_token_id", -1):
+            raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
+        self.image_token_id = int(ids)
+        preparer.image_token_id = self.image_token_id
+
+    def _encode(self, text: str) -> list[int]:
+        try:
+            ids = self.tokenizer.encode(text, add_special_tokens=False)
+        except TypeError:
+            ids = self.tokenizer.encode(text)
+        if ids and isinstance(ids[0], list):
+            ids = ids[0]
+        return list(cast("list[int]", ids))
+
+    def __call__(self, text: str, images: Sequence[Image.Image] | None = None, **_kw: Any) -> BatchFeature:
+        image_list = list(images or [])
+        token_ids = self._encode(text)
+        found = token_ids.count(self.image_token_id)
+        if found < len(image_list):
+            extra = " ".join([IMAGE_PLACEHOLDER] * (len(image_list) - found))
+            text = f"{text}\n{extra}" if found else text + extra
+            token_ids = self._encode(text)
+            found = token_ids.count(self.image_token_id)
+        if found != len(image_list):
+            raise ValueError(
+                f"Found {found} image tokens but got {len(image_list)} images"
+            )
+
+        patches: list[torch.Tensor] = []
+        grids: list[torch.Tensor] = []
+        types: list[torch.Tensor] = []
+        perms: list[torch.Tensor] = []
+        perm_counts: list[int] = []
+        running = 0
+        image_idx = 0
+        for pos, tok in enumerate(token_ids):
+            if tok != self.image_token_id:
+                running += 1
+                continue
+            prep = self.preparer.prepare(image_list[image_idx])
+            n_vit_h, n_vit_w, n_llm_h, n_llm_w = (
+                prep["n_vit_h"],
+                prep["n_vit_w"],
+                prep["n_llm_h"],
+                prep["n_llm_w"],
+            )
+            blk_types, perm = build_image_block(n_llm_h, n_llm_w, running)
+            patches.append(prep["patches"])
+            grids.append(
+                torch.tensor([n_vit_h, n_vit_w, n_llm_h, n_llm_w], dtype=torch.long)
+            )
+            types.append(blk_types)
+            perms.append(perm)
+            perm_counts.append(int(perm.numel()))
+            image_idx += 1
+            running += int(blk_types.numel())
+
+        data: dict[str, Any] = {"input_ids": [token_ids]}
+        if patches:
+            data["image_patches"] = torch.cat(patches)
+            data["image_grids"] = torch.stack(grids)
+            data["image_types"] = torch.cat(types)
+            data["image_perm"] = torch.cat(perms)
+            data["image_patch_counts"] = torch.tensor(
+                [int(pt.shape[0]) for pt in patches], dtype=torch.long
+            )
+            data["image_type_counts"] = torch.tensor(
+                [int(t.numel()) for t in types], dtype=torch.long
+            )
+            data["image_perm_counts"] = torch.tensor(perm_counts, dtype=torch.long)
+        return BatchFeature(data=data)
+
+
+class DSV4VisionDummyInputsBuilder(BaseDummyInputsBuilder[DSV4VisionProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> MultiModalDataDict:
+        num = mm_counts.get("image", 0)
+        if not num:
+            return {}
+        width, height = self.info.max_image_size
+        overrides = cast(ImageDummyOptions | None, mm_options.get("image"))
+        if overrides is not None and overrides.width and overrides.height:
+            width, height = overrides.width, overrides.height
+        image = Image.new("RGB", (width, height), color=(127, 127, 127))
+        return {"image": [image] * num}
+
+
+class DSV4VisionMultiModalProcessor(BaseMultiModalProcessor[DSV4VisionProcessingInfo]):
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        return self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**mm_kwargs),
+            dict(text=prompt, **mm_data),
+            dict(**mm_kwargs, **tok_kwargs),
+        )
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        return False
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        if "image_patches" not in hf_inputs:
+            return {}
+        patch_counts = hf_inputs["image_patch_counts"]
+        type_counts = hf_inputs["image_type_counts"]
+        perm_counts = hf_inputs["image_perm_counts"]
+        ones = torch.ones(len(patch_counts), dtype=torch.long)
+        fields: dict[str, MultiModalFieldConfig] = {
+            "image_patches": MultiModalFieldConfig.flat_from_sizes(
+                "image", patch_counts
+            ),
+            "image_perm": MultiModalFieldConfig.flat_from_sizes(
+                "image", perm_counts
+            ),
+            "image_types": MultiModalFieldConfig.flat_from_sizes(
+                "image", type_counts
+            ),
+        }
+        for key in (
+            "image_grids",
+            "image_patch_counts",
+            "image_type_counts",
+            "image_perm_counts",
+        ):
+            fields[key] = MultiModalFieldConfig.flat_from_sizes("image", ones)
+        return fields
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence:
+        token_id = self.info.processor.image_token_id
+
+        def replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
+            count = int(out_mm_kwargs["image"][item_idx]["image_type_counts"].data)
+            return PromptUpdateDetails.select_token_id(
+                [token_id] * count, token_id
+            )
+
+        return [
+            PromptReplacement(
+                modality="image",
+                target=token_id,
+                replacement=replacement,
+            )
+        ]

docker-images/dsv4-flash-vision/plugin/encoding_dsv4.py

@@ -1,957 +0,0 @@
-"""
-DeepSeek-V4 Text and Vision Encoding
-
-A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
-with tool calling, thinking mode, quick instruction tasks, and image content blocks.
-"""
-
-from typing import Any, Dict, List, Union, Optional, Tuple
-import copy
-import json
-import re
-
-# ============================================================
-# Special Tokens
-# ============================================================
-
-bos_token: str = "<｜begin▁of▁sentence｜>"
-eos_token: str = "<｜end▁of▁sentence｜>"
-thinking_start_token: str = "<think>"
-thinking_end_token: str = "</think>"
-dsml_token: str = "｜DSML｜"
-
-USER_SP_TOKEN = "<｜User｜>"
-ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
-LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
-IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"
-IMAGE_TAG_PATTERN = re.compile(r"<image>(.*?)</image>", re.DOTALL)
-
-# Task special tokens for internal classification tasks
-DS_TASK_SP_TOKENS = {
-    "action": "<｜action｜>",
-    "query": "<｜query｜>",
-    "authority": "<｜authority｜>",
-    "domain": "<｜domain｜>",
-    "title": "<｜title｜>",
-    "read_url": "<｜read_url｜>",
-}
-VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
-
-# ============================================================
-# Templates
-# ============================================================
-
-system_msg_template: str = "{content}"
-user_msg_template: str = "{content}"
-latest_reminder_msg_template: str = "{content}"
-assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
-assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
-thinking_template: str = "{reasoning_content}"
-
-response_format_template: str = (
-    "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
-)
-tool_call_template: str = (
-    "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"
-)
-tool_calls_template = (
-    "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
-)
-tool_calls_block_name: str = "tool_calls"
-
-tool_output_template: str = (
-    "<tool_result>{content}</tool_result>"
-)
-
-# Reasoning effort levels. In thinking mode, the prompt for the selected level is
-# prepended at the very beginning of the conversation. `low` is the default and
-# adds nothing.
-REASONING_EFFORT_PROMPTS: Dict[str, str] = {
-    "low": "",
-    "high": (
-        "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
-        "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
-        "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
-    ),
-    "max": (
-        "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n"
-        "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n"
-        "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n"
-    ),
-}
-DEFAULT_REASONING_EFFORT = "low"
-
-TOOLS_TEMPLATE = """## Tools
-
-You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
-
-<{dsml_token}tool_calls>
-<{dsml_token}invoke name="$TOOL_NAME">
-<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
-...
-</{dsml_token}invoke>
-<{dsml_token}invoke name="$TOOL_NAME2">
-...
-</{dsml_token}invoke>
-</{dsml_token}tool_calls>
-
-String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
-
-If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
-
-Otherwise, output directly after {thinking_end_token} with tool calls or final response.
-
-### Available Tool Schemas
-
-{tool_schemas}
-
-You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
-"""
-
-# ============================================================
-# Utility Functions
-# ============================================================
-
-def to_json(value: Any) -> str:
-    """Serialize a value to JSON string."""
-    try:
-        return json.dumps(value, ensure_ascii=False)
-    except:
-        return json.dumps(value, ensure_ascii=True)
-
-
-def tools_from_openai_format(tools):
-    """Extract function definitions from OpenAI-format tool list."""
-    return [tool["function"] for tool in tools]
-
-
-def tool_calls_from_openai_format(tool_calls):
-    """Convert OpenAI-format tool calls to internal format."""
-    return [
-        {
-            "name": tool_call["function"]["name"],
-            "arguments": tool_call["function"]["arguments"],
-        }
-        for tool_call in tool_calls
-    ]
-
-
-def tool_calls_to_openai_format(tool_calls):
-    """Convert internal tool calls to OpenAI format."""
-    return [
-        {
-            "type": "function",
-            "function": {
-                "name": tool_call["name"],
-                "arguments": tool_call["arguments"],
-            }
-        }
-        for tool_call in tool_calls
-    ]
-
-
-def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
-    """
-    Encode tool call arguments into DSML parameter format.
-
-    Args:
-        tool_call: Dict with "name" and "arguments" (JSON string) keys.
-
-    Returns:
-        DSML-formatted parameter string.
-    """
-    p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
-    P_dsml_strs = []
-
-    try:
-        arguments = json.loads(tool_call["arguments"])
-    except Exception as err:
-        arguments = {"arguments": tool_call["arguments"]}
-
-    for k, v in arguments.items():
-        p_dsml_str = p_dsml_template.format(
-            dsml_token=dsml_token,
-            key=k,
-            is_str="true" if isinstance(v, str) else "false",
-            value=v if isinstance(v, str) else to_json(v),
-        )
-        P_dsml_strs.append(p_dsml_str)
-
-    return "\n".join(P_dsml_strs)
-
-
-def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
-    """
-    Decode DSML parameters back to a tool call dict.
-
-    Args:
-        tool_name: Name of the tool.
-        tool_args: Dict mapping param_name -> (value, is_string_flag).
-
-    Returns:
-        Dict with "name" and "arguments" (JSON string) keys.
-    """
-    def _decode_value(key: str, value: str, string: str):
-        if string == "true":
-            value = to_json(value)
-        return f"{to_json(key)}: {value}"
-
-    tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
-    return dict(name=tool_name, arguments=tool_args_json)
-
-
-def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
-    """
-    Render tool schemas into the system prompt format.
-
-    Args:
-        tools: List of tool schema dicts (each with name, description, parameters).
-
-    Returns:
-        Formatted tools section string.
-    """
-    tools_json = [to_json(t) for t in tools]
-
-    return TOOLS_TEMPLATE.format(
-        tool_schemas="\n".join(tools_json),
-        dsml_token=dsml_token,
-        thinking_start_token=thinking_start_token,
-        thinking_end_token=thinking_end_token,
-    )
-
-
-def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
-    """Find the index of the last user/developer message."""
-    last_user_index = -1
-    for idx in range(len(messages) - 1, -1, -1):
-        if messages[idx].get("role") in ["user", "developer"]:
-            last_user_index = idx
-            break
-    return last_user_index
-
-
-# ============================================================
-# Message Rendering
-# ============================================================
-
-def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
-    """
-    Render a single message at the given index into its encoded string form.
-
-    This is the core function that converts each message in the conversation
-    into the DeepSeek-V4 format.
-
-    Args:
-        index: Index of the message to render.
-        messages: Full list of messages in the conversation.
-        thinking_mode: Either "chat" or "thinking".
-        drop_thinking: Whether to drop reasoning content from earlier turns.
-        reasoning_effort: Reasoning effort level, one of "low", "high", "max".
-            None is treated as "low".
-
-    Returns:
-        Encoded string for this message.
-    """
-    assert 0 <= index < len(messages)
-    assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
-
-    prompt = ""
-    msg = messages[index]
-    last_user_idx = find_last_user_index(messages)
-
-    role = msg.get("role")
-    content = msg.get("content")
-    tools = msg.get("tools")
-    response_format = msg.get("response_format")
-    tool_calls = msg.get("tool_calls")
-    reasoning_content = msg.get("reasoning_content")
-    wo_eos = msg.get("wo_eos", False)
-
-    if tools:
-        tools = tools_from_openai_format(tools)
-    if tool_calls:
-        tool_calls = tool_calls_from_openai_format(tool_calls)
-
-    # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing)
-    reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT
-    assert reasoning_effort in REASONING_EFFORT_PROMPTS, \
-        f"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}"
-    if index == 0 and thinking_mode == "thinking":
-        prompt += REASONING_EFFORT_PROMPTS[reasoning_effort]
-
-    if role == "system":
-        prompt += system_msg_template.format(content=content or "")
-        if tools:
-            prompt += "\n\n" + render_tools(tools)
-        if response_format:
-            prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
-
-    elif role == "developer":
-        assert content, f"Invalid message for role `{role}`: {msg}"
-
-        content_developer = USER_SP_TOKEN
-        content_developer += content
-
-        if tools:
-            content_developer += "\n\n" + render_tools(tools)
-        if response_format:
-            content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
-
-        prompt += user_msg_template.format(content=content_developer)
-
-    elif role == "user":
-        prompt += USER_SP_TOKEN
-
-        # Handle content blocks (tool results mixed with text)
-        content_blocks = msg.get("content_blocks")
-        if content_blocks:
-            parts = []
-            for block in content_blocks:
-                block_type = block.get("type")
-                if block_type == "text":
-                    parts.append(block.get("text", ""))
-                elif block_type == "tool_result":
-                    tool_content = block.get("content", "")
-                    if isinstance(tool_content, list):
-                        text_parts = []
-                        for b in tool_content:
-                            if b.get("type") == "text":
-                                text_parts.append(b.get("text", ""))
-                            else:
-                                text_parts.append(f"[Unsupported {b.get('type')}]")
-                        tool_content = "\n\n".join(text_parts)
-                    parts.append(tool_output_template.format(content=tool_content))
-                else:
-                    parts.append(f"[Unsupported {block_type}]")
-            prompt += "\n\n".join(parts)
-        else:
-            prompt += content or ""
-
-    elif role == "latest_reminder":
-        prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
-
-    elif role == "tool":
-        raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
-
-    elif role == "assistant":
-        thinking_part = ""
-        tc_content = ""
-
-        if tool_calls:
-            tc_list = [
-                tool_call_template.format(
-                    dsml_token=dsml_token,
-                    name=tc.get("name"),
-                    arguments=encode_arguments_to_dsml(tc)
-                )
-                for tc in tool_calls
-            ]
-            tc_content += '\n\n' + tool_calls_template.format(
-                dsml_token=dsml_token,
-                tool_calls="\n".join(tc_list),
-                tc_block_name=tool_calls_block_name,
-            )
-
-        summary_content = content or ""
-        rc = reasoning_content or ""
-
-        # Check if previous message has a task - if so, this is a task output (no thinking)
-        prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
-
-        if thinking_mode == "thinking" and not prev_has_task:
-            if not drop_thinking or index > last_user_idx:
-                thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token
-            else:
-                thinking_part = ""
-
-        if wo_eos:
-            prompt += assistant_msg_wo_eos_template.format(
-                reasoning=thinking_part,
-                content=summary_content,
-                tool_calls=tc_content,
-            )
-        else:
-            prompt += assistant_msg_template.format(
-                reasoning=thinking_part,
-                content=summary_content,
-                tool_calls=tc_content,
-            )
-    else:
-        raise NotImplementedError(f"Unknown role: {role}")
-
-    # Append transition tokens based on what follows
-    if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
-        return prompt
-
-    task = messages[index].get("task")
-    if task is not None:
-        # Task special token for internal classification tasks
-        assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
-        task_sp_token = DS_TASK_SP_TOKENS[task]
-
-        if task != "action":
-            # Non-action tasks: append task sp token directly after the message
-            prompt += task_sp_token
-        else:
-            # Action task: append Assistant + thinking token + action sp token
-            prompt += ASSISTANT_SP_TOKEN
-            prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
-            prompt += task_sp_token
-
-    elif messages[index].get("role") in ["user", "developer"]:
-        # Normal generation: append Assistant + thinking token
-        prompt += ASSISTANT_SP_TOKEN
-        if not drop_thinking and thinking_mode == "thinking":
-            prompt += thinking_start_token
-        elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
-            prompt += thinking_start_token
-        else:
-            prompt += thinking_end_token
-
-    return prompt
-
-
-# ============================================================
-# Preprocessing
-# ============================================================
-
-def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """
-    Merge tool messages into the preceding user message using content_blocks format.
-
-    DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
-    are encoded as <tool_result> blocks within user messages.
-
-    This function converts a standard OpenAI-format conversation (with separate
-    "tool" role messages) into V4 format where tool results are merged into
-    user messages.
-
-    Args:
-        messages: List of message dicts in OpenAI format.
-
-    Returns:
-        Processed message list with tool messages merged into user messages.
-    """
-    merged: List[Dict[str, Any]] = []
-
-    for msg in messages:
-        msg = copy.deepcopy(msg)
-        role = msg.get("role")
-
-        if role == "tool":
-            # Convert tool message to a user message with tool_result block
-            tool_block = {
-                "type": "tool_result",
-                "tool_use_id": msg.get("tool_call_id", ""),
-                "content": msg.get("content", ""),
-            }
-            # Merge into previous message if it's already a user (merged tool)
-            if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
-                merged[-1]["content_blocks"].append(tool_block)
-            else:
-                merged.append({
-                    "role": "user",
-                    "content_blocks": [tool_block],
-                })
-        elif role == "user":
-            content_blocks = msg.get("content_blocks")
-            if content_blocks is None:
-                content_blocks = [{"type": "text", "text": msg.get("content", "")}]
-            if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
-                merged[-1]["content_blocks"].extend(content_blocks)
-            else:
-                # Preserve structured content and all message-level metadata.
-                new_msg = msg
-                new_msg["content_blocks"] = content_blocks
-                merged.append(new_msg)
-        else:
-            merged.append(msg)
-
-    return merged
-
-
-def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """
-    Sort tool_result blocks within user messages by the order of tool_calls
-    in the preceding assistant message.
-
-    Args:
-        messages: Preprocessed message list (after merge_tool_messages).
-
-    Returns:
-        Message list with sorted tool result blocks.
-    """
-    last_tool_call_order: Dict[str, int] = {}
-
-    for msg in messages:
-        role = msg.get("role")
-        if role == "assistant" and msg.get("tool_calls"):
-            last_tool_call_order = {}
-            for idx, tc in enumerate(msg["tool_calls"]):
-                tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
-                if tc_id:
-                    last_tool_call_order[tc_id] = idx
-
-        elif role == "user" and msg.get("content_blocks"):
-            tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
-            if len(tool_blocks) > 1 and last_tool_call_order:
-                sorted_blocks = sorted(
-                    tool_blocks,
-                    key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
-                )
-                sorted_idx = 0
-                new_blocks = []
-                for block in msg["content_blocks"]:
-                    if block.get("type") == "tool_result":
-                        new_blocks.append(sorted_blocks[sorted_idx])
-                        sorted_idx += 1
-                    else:
-                        new_blocks.append(block)
-                msg["content_blocks"] = new_blocks
-
-    return messages
-
-
-# ============================================================
-# Main Encoding Function
-# ============================================================
-
-def _encode_messages_text(
-    messages: List[Dict[str, Any]],
-    thinking_mode: str,
-    context: Optional[List[Dict[str, Any]]] = None,
-    drop_thinking: bool = True,
-    add_default_bos_token: bool = True,
-    reasoning_effort: Optional[str] = None,
-) -> str:
-    """
-    Encode a list of messages into the DeepSeek-V4 prompt format.
-
-    This is the main entry point for encoding conversations. It handles:
-    - BOS token insertion
-    - Thinking mode with optional reasoning content dropping
-    - Tool message merging into user messages
-    - Multi-turn conversation context
-
-    Args:
-        messages: List of message dicts to encode.
-        thinking_mode: Either "chat" or "thinking".
-        context: Optional preceding context messages (already encoded prefix).
-        drop_thinking: If True, drop reasoning_content from earlier assistant turns
-                      (only keep reasoning for messages after the last user message).
-        add_default_bos_token: Whether to prepend BOS token at conversation start.
-        reasoning_effort: Reasoning effort level, one of "low", "high", "max".
-            Only takes effect in thinking mode. None is treated as "low".
-
-    Returns:
-        The encoded prompt string.
-    """
-    context = context if context else []
-
-    # Preprocess: merge tool messages and sort tool results
-    messages = merge_tool_messages(messages)
-    messages = sort_tool_results_by_call_order(context + messages)[len(context):]
-    if context:
-        context = merge_tool_messages(context)
-        context = sort_tool_results_by_call_order(context)
-
-    full_messages = context + messages
-
-    prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
-
-    # Resolve drop_thinking: if any message has tools defined, don't drop thinking
-    effective_drop_thinking = drop_thinking
-    if any(m.get("tools") for m in full_messages):
-        effective_drop_thinking = False
-
-    if thinking_mode == "thinking" and effective_drop_thinking:
-        full_messages = _drop_thinking_messages(full_messages)
-        # After dropping, recalculate how many messages to render
-        # (context may have shrunk too)
-        num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
-        context_len = len(full_messages) - num_to_render
-    else:
-        num_to_render = len(messages)
-        context_len = len(context)
-
-    for idx in range(num_to_render):
-        prompt += render_message(
-            idx + context_len,
-            full_messages,
-            thinking_mode=thinking_mode,
-            drop_thinking=effective_drop_thinking,
-            reasoning_effort=reasoning_effort,
-        )
-
-    return prompt
-
-
-def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """
-    Drop reasoning_content and non-essential messages before the last user message.
-
-    Behavior:
-    - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
-    - Messages at or after the last user index are always kept.
-    - Assistant messages before the last user get reasoning_content removed.
-    - Developer messages before the last user are dropped entirely.
-    """
-    last_user_idx = find_last_user_index(messages)
-    result = []
-    keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
-
-    for idx, msg in enumerate(messages):
-        role = msg.get("role")
-        if role in keep_roles or idx >= last_user_idx:
-            result.append(msg)
-        elif role == "assistant":
-            msg = copy.copy(msg)
-            msg.pop("reasoning_content", None)
-            result.append(msg)
-        # developer and other roles before last_user_idx are dropped
-
-    return result
-
-
-# ============================================================
-# Vision Message Preprocessing
-# ============================================================
-
-def parse_tagged_text(text: str) -> Union[str, List[Dict[str, Any]]]:
-    """Convert ``<image>path</image>`` text into standard content blocks."""
-    matches = list(IMAGE_TAG_PATTERN.finditer(text))
-    remaining = IMAGE_TAG_PATTERN.sub("", text)
-    if "<image>" in remaining or "</image>" in remaining:
-        raise ValueError("Malformed <image>path</image> tag")
-    if not matches:
-        return text
-
-    blocks: List[Dict[str, Any]] = []
-    cursor = 0
-    for match in matches:
-        if match.start() > cursor:
-            blocks.append({"type": "text", "text": text[cursor:match.start()]})
-        path = match.group(1)
-        if not path:
-            raise ValueError("Image path must not be empty")
-        blocks.append({
-            "type": "image_url",
-            "image_url": {"url": path},
-        })
-        cursor = match.end()
-    if cursor < len(text):
-        blocks.append({"type": "text", "text": text[cursor:]})
-    return blocks
-
-
-def _is_image_block(block: Dict[str, Any]) -> bool:
-    """Return whether a content block is an OpenAI/Anthropic/internal image."""
-    return isinstance(block, dict) and block.get("type") in ("image", "image_url")
-
-
-def _extract_image(block: Dict[str, Any]) -> Dict[str, Any]:
-    """Normalize a supported image block into an internal image record."""
-    record: Dict[str, Any] = {"type": "image"}
-    if block.get("type") == "image_url":
-        image_url = block.get("image_url")
-        if isinstance(image_url, str):
-            record["url"] = image_url
-        else:
-            record["url"] = (image_url or {}).get("url", "")
-    else:
-        for key in ("source", "url", "data"):
-            if key in block:
-                record[key] = block[key]
-    if not any(record.get(key) for key in ("source", "url", "data")):
-        raise ValueError("Image block does not contain a valid source")
-    return record
-
-
-def _process_image_blocks(
-    blocks: List[Any], image_placeholder: str = IMAGE_PLACEHOLDER
-) -> Tuple[List[Any], List[Dict[str, Any]]]:
-    """Replace image blocks and collect their records in one ordered traversal."""
-    new_blocks: List[Any] = []
-    images: List[Dict[str, Any]] = []
-    for block in blocks:
-        if not isinstance(block, dict):
-            new_blocks.append(block)
-            continue
-        if _is_image_block(block):
-            new_blocks.append({"type": "text", "text": image_placeholder})
-            images.append(_extract_image(block))
-        elif block.get("type") == "tool_result" and isinstance(block.get("content"), list):
-            block = copy.copy(block)
-            block["content"], nested_images = _process_image_blocks(
-                block["content"], image_placeholder)
-            new_blocks.append(block)
-            images.extend(nested_images)
-        elif block.get("type") == "text":
-            text = block.get("text") or ""
-            if IMAGE_PLACEHOLDER in text:
-                raise ValueError(
-                    f"Text block contains image placeholder '{IMAGE_PLACEHOLDER}': "
-                    f"'{text[:100]}'. Images should be separate content blocks."
-                )
-            new_blocks.append(block)
-        else:
-            new_blocks.append(block)
-    return new_blocks, images
-
-
-def _validate_no_image_sp_tokens(msg: Dict[str, Any]) -> None:
-    """Reject user-supplied image placeholder tokens in textual fields."""
-    content = msg.get("content")
-    if isinstance(content, str) and IMAGE_PLACEHOLDER in content:
-        raise ValueError(
-            f"Message content contains image special token '{IMAGE_PLACEHOLDER}'. "
-            "Images should be provided as image content blocks."
-        )
-    reasoning_content = msg.get("reasoning_content")
-    if isinstance(reasoning_content, str) and IMAGE_PLACEHOLDER in reasoning_content:
-        raise ValueError(
-            f"reasoning_content contains image special token '{IMAGE_PLACEHOLDER}'"
-        )
-
-
-def process_image_messages(
-    messages: List[Dict[str, Any]],
-) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
-    """Normalize image blocks and return their records in prompt order."""
-    processed: List[Dict[str, Any]] = []
-    images: List[Dict[str, Any]] = []
-    for msg in messages:
-        msg = copy.deepcopy(msg)
-        _validate_no_image_sp_tokens(msg)
-
-        if isinstance(msg.get("content"), list) and "content_blocks" not in msg:
-            msg["content_blocks"] = msg.pop("content")
-
-        if msg.get("content_blocks"):
-            msg["content_blocks"], message_images = _process_image_blocks(
-                msg["content_blocks"])
-            images.extend(message_images)
-            if not isinstance(msg.get("content"), str):
-                texts = [
-                    block.get("text", "")
-                    for block in msg["content_blocks"]
-                    if isinstance(block, dict) and block.get("type") == "text"
-                ]
-                msg["content"] = "\n\n".join(texts)
-
-        processed.append(msg)
-    return processed, images
-
-
-def encode_messages(
-    messages: List[Dict[str, Any]],
-    thinking_mode: str,
-    context: Optional[List[Dict[str, Any]]] = None,
-    drop_thinking: bool = True,
-    add_default_bos_token: bool = True,
-    reasoning_effort: Optional[str] = None,
-    return_multi_modal_data: bool = False,
-) -> Any:
-    """Encode text or multimodal messages through one canonical public entrypoint.
-
-    Text-only calls preserve the original string-returning API. When
-    return_multi_modal_data is true, the result is ``(prompt, media_data)``.
-    """
-    context = context or []
-    processed_context, _ = process_image_messages(context) if context else ([], [])
-    processed_messages, images = process_image_messages(messages)
-    prompt = _encode_messages_text(
-        processed_messages,
-        thinking_mode=thinking_mode,
-        context=processed_context if processed_context else None,
-        drop_thinking=drop_thinking,
-        add_default_bos_token=add_default_bos_token,
-        reasoning_effort=reasoning_effort,
-    )
-    if return_multi_modal_data:
-        return prompt, {"images": images}
-    return prompt
-
-
-def load_cases(input_file: str) -> List[Dict[str, Any]]:
-    """Load one or more OpenAI-format conversation cases from JSON."""
-    with open(input_file) as file:
-        data = json.load(file)
-    if isinstance(data, dict):
-        data = [data]
-    elif data and isinstance(data[0], dict) and "role" in data[0]:
-        data = [{"messages": data}]
-
-    cases = []
-    for case in data:
-        messages = copy.deepcopy(case["messages"])
-        if "tools" in case:
-            if not messages:
-                raise ValueError("A case with tools must contain at least one message")
-            messages[0]["tools"] = case["tools"]
-        cases.append({
-            "messages": messages,
-            "context": case.get("context"),
-            "thinking_mode": case.get("thinking_mode"),
-            "reasoning_effort": case.get("reasoning_effort"),
-        })
-    return cases
-
-
-def encode_case(
-    case: Dict[str, Any], thinking_mode: str
-) -> Tuple[str, List[Dict[str, Any]]]:
-    """Encode one JSON case and return its current-turn image records."""
-    prompt, media_data = encode_messages(
-        case["messages"],
-        thinking_mode=case.get("thinking_mode") or thinking_mode,
-        context=case.get("context"),
-        reasoning_effort=case.get("reasoning_effort"),
-        return_multi_modal_data=True,
-    )
-    return prompt, media_data["images"]
-
-
-# ============================================================
-# Parsing (Decoding model output)
-# ============================================================
-
-def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
-    """
-    Read text from index until one of the stop strings is found.
-
-    Returns:
-        Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
-    """
-    min_pos = len(text)
-    matched_stop = None
-
-    for s in stop:
-        pos = text.find(s, index)
-        if pos != -1 and pos < min_pos:
-            min_pos = pos
-            matched_stop = s
-
-    if matched_stop:
-        content = text[index:min_pos]
-        return min_pos + len(matched_stop), content, matched_stop
-    else:
-        content = text[index:]
-        return len(text), content, None
-
-
-def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
-    """
-    Parse DSML tool calls from text starting at the given index.
-
-    Args:
-        index: Starting position in text.
-        text: The full text to parse.
-
-    Returns:
-        Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
-        Each tool call dict has "name" and "arguments" keys.
-    """
-    tool_calls: List[Dict[str, Any]] = []
-    stop_token = None
-    tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
-
-    while index < len(text):
-        index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
-        if _ != ">\n":
-            raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'")
-
-        if stop_token == tool_calls_end_token:
-            break
-
-        if stop_token is None:
-            raise ValueError("Missing special token in tool calls")
-
-        index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
-
-        p_tool_name = re.findall(r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL)
-        if len(p_tool_name) != 1:
-            raise ValueError(f"Tool name format error: '{tool_name_content}'")
-        tool_name = p_tool_name[0]
-
-        tool_args: Dict[str, Tuple[str, str]] = {}
-        while stop_token == f"<{dsml_token}parameter":
-            index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
-
-            param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
-            if len(param_kv) != 1:
-                raise ValueError(f"Parameter format error: '{param_content}'")
-            param_name, string, param_value = param_kv[0]
-
-            if param_name in tool_args:
-                raise ValueError(f"Duplicate parameter name: '{param_name}'")
-            tool_args[param_name] = (param_value, string)
-
-            index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
-            if content != ">\n":
-                raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
-
-        tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
-        tool_calls.append(tool_call)
-
-    return index, stop_token, tool_calls
-
-
-def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
-    """
-    Parse a model completion text into a structured assistant message.
-
-    This function takes the raw text output from the model (a single assistant turn)
-    and extracts:
-    - reasoning_content (thinking block)
-    - content (summary/response)
-    - tool_calls (if any)
-
-    NOTE: This function is designed to parse only correctly formatted strings and
-    will raise ValueError for malformed output.
-
-    Args:
-        text: The raw completion text (including EOS token).
-        thinking_mode: Either "chat" or "thinking".
-
-    Returns:
-        Dict with keys: "role", "content", "reasoning_content", "tool_calls".
-        tool_calls are in OpenAI format.
-    """
-    summary_content, reasoning_content, tool_calls = "", "", []
-    index, stop_token = 0, None
-    tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
-
-    is_thinking = thinking_mode == "thinking"
-    is_tool_calling = False
-
-    if is_thinking:
-        index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
-        reasoning_content = content_delta
-        assert stop_token == thinking_end_token, "Invalid thinking format: missing </think>"
-
-    index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
-    summary_content = content_delta
-    if stop_token == tool_calls_start_token:
-        is_tool_calling = True
-    else:
-        assert stop_token == eos_token, "Invalid format: missing EOS token"
-
-    if is_tool_calling:
-        index, stop_token, tool_calls = parse_tool_calls(index, text)
-
-        index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
-        assert not tool_ends_text, "Unexpected content after tool calls"
-
-    assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end"
-
-    for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
-        assert sp_token not in summary_content and sp_token not in reasoning_content, \
-            f"Unexpected special token '{sp_token}' in content"
-
-    return {
-        "role": "assistant",
-        "content": summary_content,
-        "reasoning_content": reasoning_content,
-        "tool_calls": tool_calls_to_openai_format(tool_calls)
-    }

docker-images/dsv4-flash-vision/plugin/image_processor.py

@@ -1,188 +0,0 @@
-from __future__ import annotations
-
-import base64
-import io
-import math
-from dataclasses import dataclass
-from typing import Any
-from urllib.request import urlopen
-
-IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5)
-COMPRESS_PAD_TO = 4
-
-
-@dataclass
-class ImageInput:
-    start: int
-    patches: Any
-    n_vit_h: int
-    n_vit_w: int
-    types: Any
-    perm: Any
-
-
-def grid_tokens(best_height, best_width, patch_size, downsample_ratio):
-    """Number of LLM tokens the aligner grid occupies (N-layout, incl. row/align padding)."""
-    n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio)
-    n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio)
-    num_tokens = n_llm_h * (n_llm_w + 1) + 2
-    if n_llm_h % 2 == 1:
-        num_tokens += n_llm_w + 1
-    num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
-    return n_llm_h, n_llm_w, num_tokens
-
-
-def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
-    r = height / width
-    max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
-    max_h_float = max_w_float * r
-    if max_w_float < 1.0:
-        max_w = 1
-        max_h = (max_n_token - 2) // (max_w + 1)
-        if max_h % 2 == 1:
-            max_h -= 1
-        best_width = max_w * patch_size * downsample_ratio
-        best_height = max_h * patch_size * downsample_ratio
-    elif max_h_float < 2.0:
-        max_h = 2
-        max_w = ((max_n_token - 2) // max_h) - 1
-        assert max_w > 1
-        best_width = max_w * patch_size * downsample_ratio
-        best_height = max_h * patch_size * downsample_ratio
-    else:
-        max_w = math.floor(max_w_float)
-        max_h = math.floor(max_h_float)
-        if max_h % 2 == 1:
-            max_h -= 1
-        beta = min(max_w * patch_size * downsample_ratio / width, max_h * patch_size * downsample_ratio / height)
-        best_width = math.floor(width * beta / patch_size) * patch_size
-        best_height = math.floor(height * beta / patch_size) * patch_size
-    n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
-    return n_llm_h, n_llm_w, best_height, best_width, num_tokens
-
-
-def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):
-    max_n_token -= COMPRESS_PAD_TO - 1
-    n_llm_h, n_llm_w, num_tokens = grid_tokens(best_height, best_width, patch_size, downsample_ratio)
-    budget = max_n_token
-    while num_tokens > max_n_token:
-        n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio(
-            height, width, patch_size, downsample_ratio, budget)
-        budget -= 1
-    return n_llm_h, n_llm_w, best_height, best_width
-
-
-def load_image_bytes(record) -> bytes:
-    """Load image bytes from raw/base64 data, an Anthropic source, URL, or path."""
-    data = record.get("data")
-    if isinstance(data, bytes):
-        return data
-    if isinstance(data, str):
-        return base64.b64decode(data)
-
-    source = record.get("source")
-    if isinstance(source, dict):
-        if source.get("data") is not None:
-            return base64.b64decode(source["data"])
-        if source.get("url"):
-            return load_image_bytes({"url": source["url"]})
-
-    url = record.get("url")
-    if isinstance(url, str) and url:
-        if url.startswith("data:"):
-            header, _, payload = url.partition(",")
-            if ";base64" not in header:
-                raise ValueError(f"Unsupported data URL encoding: {header}")
-            return base64.b64decode(payload)
-        if url.startswith(("http://", "https://")):
-            with urlopen(url, timeout=30) as response:
-                return response.read()
-        with open(url, "rb") as file:
-            return file.read()
-
-    raise ValueError(f"Cannot load image from record: {list(record.keys())}")
-
-
-def load_image(record, args):
-    """Load and transform one image record into ViT patches."""
-    import numpy as np
-    import torch
-    from PIL import Image, ImageOps
-
-    p = args.vision_patch_size
-    with Image.open(io.BytesIO(load_image_bytes(record))) as source:
-        image = source.convert("RGB")
-    width, height = image.size
-    if args.vision_max_wh_ratio is not None and width > height * args.vision_max_wh_ratio:
-        width = height * args.vision_max_wh_ratio
-    if 0 < width * height < args.vision_min_pixels:
-        ratio = (args.vision_min_pixels / (width * height)) ** 0.5
-        width = int(width * ratio)
-        height = int(height * ratio)
-    best_width = math.ceil(width / p) * p
-    best_height = math.ceil(height / p) * p
-    n_llm_h, n_llm_w, best_height, best_width = safe_resize(
-        height, width, best_height, best_width, p, args.vision_downsample_ratio, args.vision_max_n_token)
-    n_vit_h, n_vit_w = best_height // p, best_width // p
-    if args.vision_max_wh_ratio is not None and image.width >= args.vision_max_wh_ratio * image.height:
-        image = image.resize((best_width, best_height))
-    else:
-        image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
-    x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
-    x = ((x - 0.5) / 0.5).to(torch.bfloat16)
-    patches = x.reshape(3, n_vit_h, p, n_vit_w, p).permute(1, 3, 0, 2, 4).reshape(n_vit_h * n_vit_w, 3, p, p)
-    return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
-
-
-def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int):
-    """Builds the N-layout token types (final order) and the aligner-row order for IMAGE slots."""
-    import torch
-
-    compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO
-    pad_h = n_llm_h % 2
-    rows = n_llm_h + pad_h
-    row_len = n_llm_w + 1
-    pad_last = rows // 2 * row_len % 2 * 2
-    types = torch.tensor(([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h), dtype=torch.int64)
-    order = torch.arange(rows * row_len).view(rows // 2, 2, row_len).transpose(1, 2).reshape(-1)
-    image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64)
-    image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(n_llm_h, n_llm_w)
-    perm = image_idx[order]
-    perm = perm[perm >= 0]
-    types = torch.cat([
-        torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64),
-        torch.tensor([IMAGE_START]),
-        types[order],
-        torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64),
-        torch.tensor([IMAGE_END]),
-    ])
-    return types, perm
-
-
-def prepare_vl_inputs(prompt, images, tokenizer, args):
-    """Expand image placeholder tokens into sentinel blocks and ImageInput values."""
-    from .encoding_dsv4 import IMAGE_PLACEHOLDER
-
-    image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
-    if image_token_id is None or image_token_id == tokenizer.unk_token_id:
-        raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
-    prompt_tokens = tokenizer.encode(prompt)
-    num_placeholders = sum(token == image_token_id for token in prompt_tokens)
-    if num_placeholders != len(images):
-        raise ValueError(
-            f"Found {num_placeholders} image tokens but got {len(images)} images")
-
-    tokens, image_inputs = [], []
-    image_iter = iter(images)
-    for tok in prompt_tokens:
-        if tok != image_token_id:
-            tokens.append(tok)
-            continue
-        patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = load_image(
-            next(image_iter), args)
-        types, perm = build_image_block(n_llm_h, n_llm_w, len(tokens))
-        image_inputs.append(ImageInput(len(tokens), patches, n_vit_h, n_vit_w, types, perm))
-        tokens += (args.vocab_size + types).tolist()
-    if not image_inputs:
-        return tokens, None
-    return tokens, image_inputs

docker-images/dsv4-flash-vision/plugin/model.py

@@ -1 +0,0 @@
-404: Not Found

docker-images/dsv4-flash-vision/plugin/vision.py

@@ -1,118 +0,0 @@
-from functools import lru_cache
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-@lru_cache(8)
-def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float):
-    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
-    hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
-    wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
-    freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq
-    freqs = freqs.flatten(1)
-    return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
-
-
-def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-    dtype = x.dtype
-    x1, x2 = x.float().chunk(2, dim=-1)
-    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        dtype = x.dtype
-        x = x.float()
-        x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
-        return (self.weight * x).to(dtype)
-
-
-class PatchEmbed(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.proj = nn.Linear(3 * args.vision_patch_size ** 2, args.vision_dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.proj(x.flatten(1))
-
-
-class Attention(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.n_heads = args.vision_n_heads
-        self.head_dim = args.vision_dim // args.vision_n_heads
-        self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim)
-        self.wo = nn.Linear(args.vision_dim, args.vision_dim)
-
-    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-        n = x.size(0)
-        q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1))
-        q = apply_rotary(q, cos, sin)
-        k = apply_rotary(k, cos, sin)
-        o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1))
-        return self.wo(o.transpose(0, 1).reshape(n, -1))
-
-
-class MLP(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False)
-        self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        gate, up = self.w1(x).chunk(2, dim=-1)
-        return self.w2(F.silu(gate) * up)
-
-
-class Block(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.norm1 = RMSNorm(args.vision_dim)
-        self.attn = Attention(args)
-        self.norm2 = RMSNorm(args.vision_dim)
-        self.mlp = MLP(args)
-
-    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), cos, sin)
-        return x + self.mlp(self.norm2(x))
-
-
-class ViT(nn.Module):
-    """DeepSeek ViT: full bidirectional attention over one image with 2D RoPE."""
-    def __init__(self, args):
-        super().__init__()
-        self.rope_dim = args.vision_dim // args.vision_n_heads // 2
-        self.rope_theta = args.vision_rope_theta
-        self.patch_embed = PatchEmbed(args)
-        self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)])
-        self.norm = RMSNorm(args.vision_dim)
-
-    def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
-        x = self.patch_embed(patches)
-        cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta)
-        for block in self.blocks:
-            x = block(x, cos, sin)
-        return self.norm(x)
-
-
-class Aligner(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.downsample_ratio = args.vision_downsample_ratio
-        in_dim = args.vision_dim * self.downsample_ratio ** 2
-        self.w1 = nn.Linear(in_dim, args.dim)
-        self.w2 = nn.Linear(args.dim, args.dim)
-
-    def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor:
-        r = self.downsample_ratio
-        x = x.view(n_h, n_w, -1).permute(2, 0, 1)
-        x = F.pad(x, (0, -n_w % r, 0, -n_h % r))
-        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
-        return self.w2(F.gelu(self.w1(x)))

src/docker/build_dir.rs

@@ -14,6 +14,54 @@ use std::path::Path;
 
 /// Compute a SHA-256 checksum over all files in `dir` (recursive).
 /// Sorted by relative path to ensure reproducibility.
+/// Copy every file under `src_dir` into `dst`: `patches/*` and `patches/`-like
+/// subtrees are flattened to `dst` root (gists cannot carry directories), other
+/// sub-directories are preserved with their relative path. Returns `true` if
+/// any non-`patches` sub-directory was preserved (callers refuse the gist flow
+/// for such overlays).
+pub fn copy_flat(src_dir: &Path, dst: &Path) -> Result<bool> {
+    let mut nested = false;
+    for entry in std::fs::read_dir(src_dir)? {
+        let entry = entry?;
+        let ft = entry.file_type()?;
+        let name = entry.file_name().to_string_lossy().to_string();
+        if ft.is_dir() {
+            if name == "patches" {
+                for f in std::fs::read_dir(entry.path())? {
+                    let f = f?;
+                    if f.file_type()?.is_file() {
+                        std::fs::copy(f.path(), dst.join(f.file_name()))
+                            .with_context(|| "copying patch")?;
+                    }
+                }
+            } else {
+                nested = true;
+                let sub = dst.join(&name);
+                std::fs::create_dir_all(&sub)?;
+                copy_tree(&entry.path(), &sub)?;
+            }
+        } else if ft.is_file() && name != "build.sh" && name != "Dockerfile" {
+            std::fs::copy(entry.path(), dst.join(&name))?;
+        }
+    }
+    Ok(nested)
+}
+
+fn copy_tree(src_dir: &Path, dst_dir: &Path) -> Result<()> {
+    for entry in std::fs::read_dir(src_dir)? {
+        let entry = entry?;
+        let ft = entry.file_type()?;
+        let target = dst_dir.join(entry.file_name());
+        if ft.is_dir() {
+            std::fs::create_dir_all(&target)?;
+            copy_tree(&entry.path(), &target)?;
+        } else if ft.is_file() {
+            std::fs::copy(entry.path(), &target)?;
+        }
+    }
+    Ok(())
+}
+
 pub fn compute_dir_checksum(dir: &Path) -> Result<String> {
     let mut hasher = Sha256::new();
     let mut entries: Vec<_> = Vec::new();

src/publish.rs

@@ -3,8 +3,9 @@
 //! patches, run/install scripts, the recipe itself). `--gist` additionally
 //! publishes the bundle as a public GitHub gist via `gh`.
 //!
-//! The bundle is intentionally flat (gists have no directories): `patches/`
-//! paths in the Dockerfile are rewritten to the bundle root.
+//! `patches/` files are flattened into the bundle root (gists cannot carry
+//! directories); other sub-directories are preserved and `--gist` is refused
+//! for such overlays.
 
 use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
@@ -14,6 +15,7 @@ use anyhow::{Context, Result, bail};
 use tracing::info;
 
 use crate::config::{Cluster, Config, NcclConfig, Recipe, recipe_dirs};
+use crate::docker::build_dir::copy_flat;
 use crate::ssh::shq;
 
 /// Publish `recipe_name` into `out` (default `./publish/<name>`), optionally
@@ -52,7 +54,12 @@ pub async fn publish(
     if image_dir.join("Dockerfile").exists() {
         let df = std::fs::read_to_string(image_dir.join("Dockerfile"))?;
         std::fs::write(bundle.join("Dockerfile"), flatten_dockerfile(&df))?;
-        copy_flat(&image_dir, &bundle)?;
+        let nested = copy_flat(&image_dir, &bundle)?;
+        if gist && nested {
+            bail!(
+                "publish --gist needs a flat overlay ({image_dir:?} nests non-patches dirs) — use --out/--registry instead"
+            );
+        }
         if let Ok(v) = std::fs::read_to_string(image_dir.join("versions.env")) {
             std::fs::write(bundle.join("versions.env"), &v)?;
             versions = versions_build_args(&v);
@@ -178,35 +185,6 @@ pub(crate) fn versions_build_args(text: &str) -> BTreeMap<String, String> {
     out
 }
 
-/// Copy every file under `dir` (including `patches/*`) into `dst` root,
-/// flattening paths. Sub-directories other than `patches/` are rejected —
-/// a flat bundle can only reference root-level files.
-fn copy_flat(src_dir: &Path, dst: &Path) -> Result<()> {
-    for entry in std::fs::read_dir(src_dir)? {
-        let entry = entry?;
-        let ft = entry.file_type()?;
-        let name = entry.file_name().to_string_lossy().to_string();
-        if ft.is_dir() {
-            if name != "patches" {
-                bail!(
-                    "publish: unsupported sub-dir '{}' in image overlay (flat bundles only)",
-                    name
-                );
-            }
-            for f in std::fs::read_dir(entry.path())? {
-                let f = f?;
-                if f.file_type()?.is_file() {
-                    std::fs::copy(f.path(), dst.join(f.file_name()))
-                        .with_context(|| "copying patch")?;
-                }
-            }
-        } else if ft.is_file() && name != "build.sh" && name != "Dockerfile" {
-            std::fs::copy(entry.path(), dst.join(&name))?;
-        }
-    }
-    Ok(())
-}
-
 fn render_run_script(
     recipe: &Recipe,
     cluster: Option<&Cluster>,
@@ -463,6 +441,27 @@ mod tests {
         assert!(f.contains("  c.patch\nRUN grep patches/"));
     }
 
+    #[test]
+    fn copy_flat_preserves_nested_dirs_reports_nested() {
+        let tmp = std::env::temp_dir().join("sm121_copy_flat_test");
+        let _ = std::fs::remove_dir_all(&tmp);
+        let src = tmp.join("src");
+        std::fs::create_dir_all(src.join("plugin/pkg")).unwrap();
+        std::fs::create_dir_all(src.join("patches")).unwrap();
+        std::fs::write(src.join("plugin/pyproject.toml"), "x").unwrap();
+        std::fs::write(src.join("plugin/pkg/m.py"), "y").unwrap();
+        std::fs::write(src.join("patches/p.patch"), "z").unwrap();
+        std::fs::write(src.join("Dockerfile"), "f").unwrap();
+        let dst = tmp.join("dst");
+        std::fs::create_dir_all(&dst).unwrap();
+        let nested = copy_flat(&src, &dst).unwrap();
+        assert!(nested);
+        let m = std::fs::read_to_string(dst.join("plugin/pkg/m.py")).unwrap();
+        assert_eq!(m, "y");
+        assert_eq!(std::fs::read_to_string(dst.join("p.patch")).unwrap(), "z");
+        assert!(!dst.join("Dockerfile").exists());
+    }
+
     #[test]
     fn versions_env_parses_quoted_and_commented() {
         let v = versions_build_args("# c\nBASE_IMAGE=\"img@sha256:1\"\nX = 'y' \nNOVAL=\n");