bundle/template.py

"""Lazy explicit registration for the frozen Qwen3.5 realtime template."""

from __future__ import annotations

import importlib
from typing import Any

TEMPLATE_NAME = "qwen3_5_realtime"
FROZEN_SOURCE_COMMIT = "2132d5e989d8f7dbaed8abfb78268dafbd6f5402"
FROZEN_PLUGIN_PATH = (
    "procedure_harness/research/timed_video_continuation_b/realtime_template_plugin.py"
)
FROZEN_PLUGIN_SHA256 = (
    "e8a51bc12adbc11898115c1c6741c17a2c0c5a8cc43264e7aae7916aff5d1290"
)

_REGISTERED_TEMPLATE_CLASS: type[Any] | None = None


def frozen_runtime_config() -> dict[str, Any]:
    """Return a fresh copy of the frozen launch and generation identity."""

    return {
        "template": TEMPLATE_NAME,
        "max_pixels": 147456,
        "max_model_len": 32768,
        "max_tokens": 128,
        "temperature": 0.0,
        "engine_seed": 43,
        "sampling_seed": 43,
        "tensor_parallel_size": 1,
        "enable_thinking": False,
        "add_non_thinking_prefix": True,
    }


def _registration_receipt(status: str) -> dict[str, str]:
    return {
        "status": status,
        "template": TEMPLATE_NAME,
        "source_commit": FROZEN_SOURCE_COMMIT,
        "source_path": FROZEN_PLUGIN_PATH,
        "source_sha256": FROZEN_PLUGIN_SHA256,
    }


def register_realtime_template() -> dict[str, str]:
    """Explicitly register the vendored template and return its provenance.

    Heavy runtime packages are imported only inside this call. A conflicting
    pre-existing registration is rejected instead of silently overriding the
    process-wide ms-swift template mapping.
    """

    global _REGISTERED_TEMPLATE_CLASS

    register_module = importlib.import_module("swift.template.register")
    mapping = register_module.TEMPLATE_MAPPING
    existing = mapping.get(TEMPLATE_NAME)
    if existing is not None:
        if (
            _REGISTERED_TEMPLATE_CLASS is None
            or getattr(existing, "template_cls", None) is not _REGISTERED_TEMPLATE_CLASS
        ):
            raise RuntimeError(
                f"refusing conflicting {TEMPLATE_NAME!r} template registration"
            )
        return _registration_receipt("already_registered")

    torch = importlib.import_module("torch")
    qwen_module = importlib.import_module("swift.template.templates.qwen")
    qwen_base = qwen_module.Qwen3_5Template
    qwen_meta = qwen_module.QwenTemplateMeta

    class Qwen3_5RealTimeTemplate(qwen_base):
        """Qwen3.5 template that replaces fake FPS with absolute source PTS."""

        frozen_source_commit = FROZEN_SOURCE_COMMIT
        frozen_source_sha256 = FROZEN_PLUGIN_SHA256

        def replace_tag(self, media_type: str, index: int, inputs: Any) -> list[Any]:
            qwen_utils = importlib.import_module("qwen_vl_utils")
            fetch_image = qwen_utils.fetch_image
            fetch_video = qwen_utils.fetch_video

            if media_type not in {"image", "video"}:
                raise AssertionError(f"unsupported media type: {media_type}")
            kwargs = (
                {"image_patch_size": self.processor.image_processor.patch_size}
                if self.version == "v3"
                else {}
            )
            if self.mode == "vllm":
                inputs.mm_processor_kwargs["do_resize"] = False
            if media_type == "image":
                inputs.images[index] = fetch_image(
                    {"image": inputs.images[index], **inputs.chat_template_kwargs},
                    **kwargs,
                )
                if self.mode == "lmdeploy":
                    return ["<|vision_start|>", [-100], "<|vision_end|>"]
                return ["<|vision_start|><|image_pad|><|vision_end|>"]

            if self.version == "v3":
                kwargs["return_video_metadata"] = True
            video = inputs.videos[index]
            timed_video = None
            if (
                isinstance(video, dict)
                and video.get("schema_version") == "timed_video.v1"
            ):
                if self.version != "v3":
                    raise RuntimeError("timed_video.v1 requires Qwen v3 metadata")
                from .realtime import TimedVideo

                timed_video = TimedVideo.from_mapping(video)
                video = list(timed_video.processor_view().frame_paths)
            video_inputs = {"video": video, **inputs.chat_template_kwargs}
            if isinstance(video, list):
                vision_process = importlib.import_module("qwen_vl_utils.vision_process")
                video_inputs["sample_fps"] = vision_process.FPS
            video, _video_kwargs = fetch_video(
                video_inputs,
                return_video_sample_fps=True,
                **kwargs,
            )
            tokens: list[Any] = ["<|vision_start|><|video_pad|><|vision_end|>"]
            if self.version == "v2_5":
                inputs.mm_processor_kwargs.setdefault("fps", []).append(_video_kwargs)
            elif self.version == "v3":
                if self.mode != "vllm":
                    video, video_metadata = video
                    if timed_video is not None:
                        video_metadata = timed_video.video_metadata()
                    inputs.mm_processor_kwargs.setdefault("video_metadata", []).append(
                        video_metadata
                    )
                    tokens = ["<|video_pad|>"]
                elif timed_video is not None:
                    decoded_video, _video_metadata = video
                    video = (decoded_video, timed_video.video_metadata())
                inputs.mm_processor_kwargs["do_sample_frames"] = False
            if isinstance(video, torch.Tensor):
                video = video.to(torch.uint8)
            inputs.videos[index] = video
            return tokens

    Qwen3_5RealTimeTemplate.__module__ = __name__
    Qwen3_5RealTimeTemplate.__qualname__ = "Qwen3_5RealTimeTemplate"
    meta = qwen_meta(
        TEMPLATE_NAME,
        template_cls=Qwen3_5RealTimeTemplate,
        default_system=None,
        thinking_prefix="<think>\n",
        non_thinking_prefix="<think>\n\n</think>\n\n",
        agent_template="qwen3_5",
        is_thinking=True,
    )
    register_module.register_template(meta)
    registered = mapping.get(TEMPLATE_NAME)
    if getattr(registered, "template_cls", None) is not Qwen3_5RealTimeTemplate:
        raise RuntimeError(f"failed to register {TEMPLATE_NAME!r} template")
    _REGISTERED_TEMPLATE_CLASS = Qwen3_5RealTimeTemplate
    return _registration_receipt("registered")