from transformers import AutoConfig, AutoProcessor, StoppingCriteria
from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
    Qwen3OmniMoeThinkerForConditionalGeneration,
)
from collections import OrderedDict

import os
import torch
import torch.nn.functional as F

from toolkit.basic import flush
from toolkit.util.comfy_quant_import import (
    import_comfy_quantized_layers,
    parse_comfy_quant_blob,
)
from toolkit.util.convrot_quant import regular_hadamard

from .BaseCaptioner import BaseCaptioner
from .Qwen3VLCaptioner import patch_qwen_vl_patch_embed
import logging
import traceback
import warnings

warnings.filterwarnings("ignore")
logging.disable(logging.WARNING)

# frame sampling rate for video captioning
VIDEO_FPS = 2

# still-image files caption through the image pipeline (no audio, no frames)
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"}
# audio-only files caption through the audio pipeline (no frames)
AUDIO_EXTENSIONS = {"mp3", "wav", "flac", "ogg", "m4a", "aac"}

# fixed generation ceiling under compiled decode: a constant max_length keeps
# the static kv cache (and so the compiled decode graph) at one shape for
# every video; the real per-caption budget is enforced by a stopping criterion
STATIC_MAX_LENGTH = 8192

# reasoning cap for thinking models: the visible caption gets the full
# max_new_tokens budget only after </think> closes
MAX_THINKING_TOKENS = 4096

# single-file comfy-format checkpoints (thinker only, convrot8 int8) produced
# by scripts/convert_vllm_to_comfy.py. This is always what we load — never the
# original bf16 shards. base_repo supplies config + processor (tokenizer,
# feature extractors, chat template — thinking models need the thinking
# template, which the finetune repos don't always ship).
CONVROT_MODELS = {
    "ai-toolkit/Qwen3-Omni-30B-A3B-Instruct": {
        "filename": "qwen3_omni_30b_a3b_instruct_thinker_convrot8.safetensors",
        "base_repo": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
        "thinking": False,
    },
    "ai-toolkit/Qwen3-Omni-30B-A3B-Thinking": {
        "filename": "qwen3_omni_30b_a3b_thinking_convrot8.safetensors",
        "base_repo": "Qwen/Qwen3-Omni-30B-A3B-Thinking",
        "thinking": True,
    },
    "ai-toolkit/Huihui-Qwen3-Omni-30B-A3B-Thinking-abliterated": {
        "filename": "huihui_qwen3_omni_30b_a3b_thinking_abliterated_convrot8.safetensors",
        "base_repo": "Qwen/Qwen3-Omni-30B-A3B-Thinking",
        "thinking": True,
    },
}
DEFAULT_CONVROT_MODEL = "ai-toolkit/Qwen3-Omni-30B-A3B-Instruct"


class BatchThinkingBudgetCriteria(StoppingCriteria):
    """Per-row thinking budget: let each sequence reason freely, then count
    max_new_tokens from the token after its </think> so the visible caption
    gets the full budget regardless of how long the reasoning ran. Rows that
    never close their think block are bounded by the accompanying
    MaxLengthCriteria / max_new_tokens ceiling."""

    def __init__(self, think_end_token_id: int, max_new_tokens: int):
        self.think_end_token_id = think_end_token_id
        self.max_new_tokens = max_new_tokens
        self.answer_start = None

    def __call__(self, input_ids, scores, **kwargs):
        batch, length = input_ids.shape
        if self.answer_start is None:
            self.answer_start = torch.full(
                (batch,), -1, dtype=torch.long, device=input_ids.device
            )
        newly_closed = (input_ids[:, -1] == self.think_end_token_id) & (
            self.answer_start < 0
        )
        self.answer_start[newly_closed] = length
        return (self.answer_start >= 0) & (
            length - self.answer_start >= self.max_new_tokens
        )


class OstrisQwen3OmniThinker(Qwen3OmniMoeThinkerForConditionalGeneration):
    """Thinker with static-cache-safe MRoPE handling.

    Upstream breaks under ``cache_implementation="static"``: generate passes a
    prepared 4D bool attention mask, but the forward's rope-delta block does
    ``1 - attention_mask`` and ``get_rope_index`` assumes a 2D long padding
    mask. We compute position_ids ourselves — prefill from the true 2D mask
    (stashed by the caller before generate), decode from cache_position with
    no data-dependent ops — so the upstream block (which only runs when
    position_ids is None) is skipped entirely. Also required for CUDA-graph
    decode: the decode branch is sync-free and shape-static."""

    _pad_mask_2d = None

    # media inputs are consumed at prefill only; keeping them in decode-step
    # inputs makes the compiled decode graph guard on their (per-video) shapes,
    # forcing a recompile on the next video. Dropping them gives the decode
    # graph one fixed signature: it compiles once, ever.
    _PREFILL_ONLY_KEYS = (
        "input_features",
        "feature_attention_mask",
        "audio_feature_lengths",
        "pixel_values",
        "pixel_values_videos",
        "image_grid_thw",
        "video_grid_thw",
        "video_second_per_grid",
    )

    def prepare_inputs_for_generation(self, *args, **kwargs):
        model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
        ids = model_inputs.get("input_ids", None)
        if ids is not None and ids.shape[1] == 1:
            for key in self._PREFILL_ONLY_KEYS:
                model_inputs.pop(key, None)
        return model_inputs

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        cache_position=None,
        input_features=None,
        pixel_values=None,
        pixel_values_videos=None,
        image_grid_thw=None,
        video_grid_thw=None,
        feature_attention_mask=None,
        audio_feature_lengths=None,
        use_audio_in_video=None,
        video_second_per_grid=None,
        **kwargs,
    ):
        if position_ids is None and input_ids is not None:
            if input_ids.shape[1] > 1 or self.rope_deltas is None:
                # prefill: replicate the upstream math with a valid 2D mask
                mask2d = (
                    attention_mask
                    if attention_mask is not None and attention_mask.dim() == 2
                    else self._pad_mask_2d
                )
                if mask2d is None:
                    mask2d = torch.ones_like(input_ids)
                mask2d = mask2d.long()
                if mask2d.shape[1] != input_ids.shape[1]:
                    # static cache pads the mask out to max_cache_len
                    mask2d = mask2d[:, : input_ids.shape[1]]
                if feature_attention_mask is not None:
                    rope_audio_lengths = torch.sum(feature_attention_mask, dim=1)
                else:
                    rope_audio_lengths = audio_feature_lengths
                delta0 = (1 - mask2d).sum(dim=-1).unsqueeze(1)
                position_ids, rope_deltas = self.get_rope_index(
                    input_ids,
                    image_grid_thw,
                    video_grid_thw,
                    mask2d,
                    use_audio_in_video or False,
                    rope_audio_lengths,
                    video_second_per_grid,
                )
                self.rope_deltas = rope_deltas - delta0
            else:
                # decode: continue from the cache position; sync-free
                batch_size, seq_length = input_ids.shape
                deltas = self.rope_deltas.to(input_ids.device)
                if cache_position is not None:
                    pos = cache_position.view(1, -1) + deltas
                else:
                    # get_seq_length may be a tensor (static cache); keep it on-device
                    past_len = (
                        past_key_values.get_seq_length()
                        if past_key_values is not None
                        else 0
                    )
                    pos = (
                        torch.arange(seq_length, device=input_ids.device).view(1, -1)
                        + past_len
                        + deltas
                    )
                position_ids = pos.unsqueeze(0).expand(3, batch_size, seq_length)
        return super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            cache_position=cache_position,
            input_features=input_features,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            feature_attention_mask=feature_attention_mask,
            audio_feature_lengths=audio_feature_lengths,
            use_audio_in_video=use_audio_in_video,
            video_second_per_grid=video_second_per_grid,
            **kwargs,
        )


class ConvRot8Experts(torch.nn.Module):
    """Drop-in replacement for Qwen3OmniMoeThinkerTextExperts that keeps the
    fused expert banks in comfy convrot8 storage (regular-Hadamard rotated,
    per-output-row symmetric int8). Experts are dequantized one at a time at
    forward, so the full-precision banks (the bulk of the 30B) never
    materialize."""

    def __init__(
        self, gate_up_q, gate_up_s, gate_up_rot, down_q, down_s, down_rot, dtype
    ):
        super().__init__()
        self.num_experts = gate_up_q.shape[0]
        self.gate_up_rot = gate_up_rot
        self.down_rot = down_rot
        self.out_dtype = dtype
        self.register_buffer("gate_up_q", gate_up_q.contiguous(), persistent=False)
        self.register_buffer("down_q", down_q.contiguous(), persistent=False)
        # fp32 scales stored as uint8 byte views so a later .to(dtype=...) on the
        # model cannot silently cast them (same convention as the cr8 backend)
        self.register_buffer(
            "gate_up_s",
            gate_up_s.detach().float().contiguous().view(torch.uint8),
            persistent=False,
        )
        self.register_buffer(
            "down_s",
            down_s.detach().float().contiguous().view(torch.uint8),
            persistent=False,
        )
        # hadamard matrices as buffers: the toolkit's cached builder is a
        # global-dict lookup that torch.compile cannot trace
        self.register_buffer(
            "gate_up_h",
            regular_hadamard(gate_up_rot, torch.device("cpu"), torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "down_h",
            regular_hadamard(down_rot, torch.device("cpu"), torch.float32),
            persistent=False,
        )

    # device the streamed experts should land on when the banks themselves
    # stay in system RAM (low-vram layer offloading); None = banks resident
    offload_device = None

    def enable_offload(self, device):
        """Keep the int8 banks in (pinned) system RAM; forward streams only
        the routed experts' rows to the GPU per layer call."""
        self.offload_device = device
        try:
            self.gate_up_q = self.gate_up_q.pin_memory()
            self.down_q = self.down_q.pin_memory()
            self.gate_up_s = self.gate_up_s.pin_memory()
            self.down_s = self.down_s.pin_memory()
        except RuntimeError:
            pass  # pinning is a speed optimization only; pageable still works
        # the hadamard matrices are tiny — keep them resident
        self.gate_up_h = self.gate_up_h.to(device)
        self.down_h = self.down_h.to(device)

    @staticmethod
    def _rotate(w, h, rot):
        shape = w.shape
        return (w.reshape(-1, shape[-1] // rot, rot) @ h).reshape(shape)

    def _gather(self, qdata, scales_u8, hit):
        """Expert rows + scales for the hit indices, on the compute device."""
        if self.offload_device is not None and qdata.device.type == "cpu":
            # each expert's rows are a contiguous view of the pinned bank, so
            # slice-copies DMA straight to the GPU with zero CPU-side gather
            # work (a CPU index_select here memcpy'd ~2GB/token on all cores)
            hit_list = hit.tolist() if torch.is_tensor(hit) else list(hit)
            scales = scales_u8.view(torch.float32)
            q = torch.stack(
                [qdata[i].to(self.offload_device, non_blocking=True) for i in hit_list]
            )
            s = torch.stack(
                [scales[i].to(self.offload_device, non_blocking=True) for i in hit_list]
            )
            return q, s
        return qdata[hit], scales_u8.view(torch.float32)[hit]

    def _dequant(self, qdata, scales_u8, h, rot, i):
        # scales are [E, out, 1]; rotation is self-inverse along the in dim
        q, s = self._gather(
            qdata, scales_u8, i.reshape(1) if torch.is_tensor(i) else torch.tensor([i])
        )
        w = q[0].float() * s[0]
        return self._rotate(w, h, rot).to(self.out_dtype)

    def _dequant_batch(self, qdata, scales_u8, h, rot, hit, dtype):
        """Dequantize the hit experts in one shot: [n_hit, out, in]."""
        q, s = self._gather(qdata, scales_u8, hit)
        w = q.float() * s
        return self._rotate(w, h, rot).to(dtype)

    def forward(self, hidden_states, top_k_index, top_k_weights):
        """Fully batched MoE: group tokens by expert (sort + bincount), pad the
        groups to a rectangle, dequantize the hit experts in one op, and run the
        whole layer as two bmms — no per-expert python loop. Decode touches only
        the routed experts' weights; prefill runs every expert in one launch."""
        hidden_dim = hidden_states.shape[1]
        top_k = top_k_index.shape[-1]

        # gate on token count, not pair count: decode (1 token per sequence)
        # must ALWAYS take this path at any batch size — the grouped path's
        # nonzero()/max() are data-dependent, and inside the compiled decode
        # graph they shatter it into per-layer fragments (endless compiles,
        # broken cudagraphs). Extra cost is only duplicate expert dequants
        # (~1.6x traffic at batch 16). Prefill (many tokens, runs eager)
        # still uses the grouped path below.
        if hidden_states.shape[0] <= 32:
            # decode-size batches: one bmm per (token, expert) pair with fixed
            # shapes and NO data-dependent ops — the grouped path below needs
            # nonzero()/max() which each force a GPU sync, and 2 syncs x 48
            # layers per token is exactly what stalls the GPU at small batch
            flat = top_k_index.reshape(-1)
            x_rep = hidden_states.repeat_interleave(top_k, dim=0).unsqueeze(1)
            w_gate_up = self._dequant_batch(
                self.gate_up_q,
                self.gate_up_s,
                self.gate_up_h,
                self.gate_up_rot,
                flat,
                hidden_states.dtype,
            )
            gate, up = torch.bmm(x_rep, w_gate_up.transpose(1, 2)).chunk(2, dim=-1)
            del w_gate_up
            h = F.silu(gate) * up
            w_down = self._dequant_batch(
                self.down_q,
                self.down_s,
                self.down_h,
                self.down_rot,
                flat,
                hidden_states.dtype,
            )
            out = torch.bmm(h, w_down.transpose(1, 2)).squeeze(1)
            del w_down
            out = out * top_k_weights.reshape(-1, 1)
            return (
                out.view(hidden_states.shape[0], top_k, hidden_dim)
                .sum(dim=1)
                .to(hidden_states.dtype)
            )
        device = hidden_states.device
        dtype = hidden_states.dtype

        flat_expert = top_k_index.reshape(-1)  # [n_tokens * top_k]
        order = flat_expert.argsort()
        sorted_expert = flat_expert[order]
        token_of_pair = order // top_k
        counts = torch.bincount(flat_expert, minlength=self.num_experts)
        hit = counts.nonzero().flatten()
        hit_counts = counts[hit]
        group_size = int(hit_counts.max())
        # rank of each routed pair inside its expert group
        group_start = (torch.cumsum(counts, 0) - counts)[sorted_expert]
        rank = torch.arange(order.shape[0], device=device) - group_start
        slot = torch.searchsorted(hit, sorted_expert)

        padded_x = torch.zeros(
            hit.shape[0], group_size, hidden_dim, device=device, dtype=dtype
        )
        padded_x[slot, rank] = hidden_states[token_of_pair]

        w_gate_up = self._dequant_batch(
            self.gate_up_q, self.gate_up_s, self.gate_up_h, self.gate_up_rot, hit, dtype
        )
        gate, up = torch.bmm(padded_x, w_gate_up.transpose(1, 2)).chunk(2, dim=-1)
        del w_gate_up
        h = F.silu(gate) * up
        w_down = self._dequant_batch(
            self.down_q, self.down_s, self.down_h, self.down_rot, hit, dtype
        )
        out = torch.bmm(h, w_down.transpose(1, 2))
        del w_down

        pair_out = out[slot, rank] * top_k_weights.reshape(-1)[order].unsqueeze(1)
        final_hidden_states = torch.zeros_like(hidden_states)
        final_hidden_states.index_add_(0, token_of_pair, pair_out.to(dtype))
        return final_hidden_states

    def _forward_dequant(self, hidden_states, top_k_index, top_k_weights):
        # mirrors Qwen3OmniMoeThinkerTextExperts.forward with per-expert dequant
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            w_gate_up = self._dequant(
                self.gate_up_q,
                self.gate_up_s,
                self.gate_up_h,
                self.gate_up_rot,
                expert_idx,
            )
            gate, up = F.linear(current_state, w_gate_up).chunk(2, dim=-1)
            current_hidden_states = F.silu(gate) * up
            w_down = self._dequant(
                self.down_q, self.down_s, self.down_h, self.down_rot, expert_idx
            )
            current_hidden_states = F.linear(current_hidden_states, w_down)
            current_hidden_states = (
                current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            )
            final_hidden_states.index_add_(
                0, token_idx, current_hidden_states.to(final_hidden_states.dtype)
            )

        return final_hidden_states


def swap_convrot_expert_banks(root, state_dict, dtype):
    """Replace each MoE experts module with a ConvRot8Experts holding the
    quantized banks from the checkpoint, consuming their state dict entries.
    Returns (remaining_state_dict, num_swapped)."""
    state_dict = dict(state_dict)
    bank_paths = sorted(
        {
            k[: -len(".gate_up_proj.comfy_quant")]
            for k in state_dict
            if k.endswith(".gate_up_proj.comfy_quant") and ".experts" in k
        }
    )
    for experts_path in bank_paths:
        tensors = {}
        rots = {}
        for proj in ("gate_up_proj", "down_proj"):
            prefix = f"{experts_path}.{proj}"
            conf = parse_comfy_quant_blob(state_dict.pop(f"{prefix}.comfy_quant"))
            if conf.get("format") != "int8_tensorwise" or not conf.get("convrot"):
                raise ValueError(
                    f"Expert bank {prefix} has unsupported quant config {conf}"
                )
            tensors[proj + "_q"] = state_dict.pop(f"{prefix}.weight")
            tensors[proj + "_s"] = state_dict.pop(f"{prefix}.weight_scale")
            rots[proj] = int(conf.get("convrot_groupsize", 256))

        parent_path, _, attr = experts_path.rpartition(".")
        parent = root.get_submodule(parent_path)
        setattr(
            parent,
            attr,
            ConvRot8Experts(
                tensors["gate_up_proj_q"],
                tensors["gate_up_proj_s"],
                rots["gate_up_proj"],
                tensors["down_proj_q"],
                tensors["down_proj_s"],
                rots["down_proj"],
                dtype,
            ),
        )
    return state_dict, len(bank_paths)


class Qwen3OmniCaptioner(BaseCaptioner):
    """Captions videos using their audio track via the Qwen3-Omni thinker,
    loaded from the pre-quantized convrot8 single-file checkpoint."""

    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
        super(Qwen3OmniCaptioner, self).__init__(process_id, job, config, **kwargs)

    def _resolve_checkpoint(self) -> str:
        """model_name_or_path can be the checkpoint file itself, a folder
        holding it, or a hub repo. Known local spots under MODELS_PATH
        (text_encoders/, the root, then any subfolder of text_encoders/) are
        searched before downloading; downloads land in
        MODELS_PATH/text_encoders."""
        from toolkit.paths import MODELS_PATH

        def info_for_filename(filename):
            for info in CONVROT_MODELS.values():
                if info["filename"] == filename:
                    return info
            return CONVROT_MODELS[DEFAULT_CONVROT_MODEL]

        name_or_path = self.caption_config.model_name_or_path
        if os.path.isfile(name_or_path):
            self._model_info = info_for_filename(os.path.basename(name_or_path))
            return name_or_path

        model_info = CONVROT_MODELS.get(
            name_or_path, CONVROT_MODELS[DEFAULT_CONVROT_MODEL]
        )
        filename = model_info["filename"]

        if os.path.isdir(name_or_path):
            candidate = os.path.join(name_or_path, filename)
            if os.path.exists(candidate):
                self._model_info = model_info
                return candidate
            files = [f for f in os.listdir(name_or_path) if f.endswith(".safetensors")]
            if len(files) == 1:
                self._model_info = info_for_filename(files[0])
                return os.path.join(name_or_path, files[0])
            raise FileNotFoundError(
                f"No {filename} (or single .safetensors) in {name_or_path}"
            )

        self._model_info = model_info
        te_dir = os.path.join(MODELS_PATH, "text_encoders")
        for candidate in (
            os.path.join(te_dir, filename),
            os.path.join(MODELS_PATH, filename),
        ):
            if os.path.exists(candidate):
                return candidate
        if os.path.isdir(te_dir):
            for dirpath, dirnames, filenames in os.walk(te_dir):
                dirnames.sort()
                if filename in filenames:
                    return os.path.join(dirpath, filename)

        import huggingface_hub

        self.print_and_status_update(
            f"Downloading {filename} from {name_or_path} into {te_dir}"
        )
        return huggingface_hub.hf_hub_download(
            repo_id=name_or_path, filename=filename, local_dir=te_dir
        )

    def load_model(self):
        from accelerate import init_empty_weights
        from safetensors.torch import load_file

        ckpt_path = self._resolve_checkpoint()
        base_repo = self._model_info["base_repo"]
        self.is_thinking_model = self._model_info["thinking"]
        # thinking models reason by default; the template's enable_thinking=False
        # (an empty <think></think> block) suppresses it unless the user asked
        self.thinking_enabled = self.is_thinking_model and self.caption_config.thinking
        self.print_and_status_update(
            f"Loading Qwen3-Omni thinker (convrot8, base {base_repo})"
        )

        config = AutoConfig.from_pretrained(base_repo)
        with init_empty_weights(include_buffers=False):
            model = OstrisQwen3OmniThinker(config.thinker_config)
        model.eval()

        # NOTE: flash_attention_2 was tried here and produced degenerate
        # repetitive output on real jobs (likely its padding handling against
        # the fixed-size static cache with left-padded batches); sdpa is
        # correct and nearly as fast, so we stay on it.

        state_dict = load_file(ckpt_path)

        # MoE expert banks stay int8 in ConvRot8Experts modules
        state_dict, num_banks = swap_convrot_expert_banks(
            model, state_dict, self.torch_dtype
        )
        # everything else quantized (attention, vision, audio linears) attaches
        # to the toolkit's convrot8 backend in place — no dequantization
        state_dict, num_quantized = import_comfy_quantized_layers(
            model, state_dict, orig_dtype=self.torch_dtype
        )
        self.print_and_status_update(
            f" - attached {num_banks} expert banks and {num_quantized} ConvRot layers"
        )
        result = model.load_state_dict(state_dict, assign=True, strict=False)
        # the importer already attached weights (and popped + assigned biases)
        # of quantized layers, so load_state_dict reports them as missing
        expected_missing = set()
        for name, module in model.named_modules():
            if hasattr(module, "ostris_quantizer"):
                expected_missing.add(f"{name}.weight")
                expected_missing.add(f"{name}.bias")
        bad_missing = [k for k in result.missing_keys if k not in expected_missing]
        if bad_missing or result.unexpected_keys:
            raise RuntimeError(
                f"Checkpoint mismatch. missing: {bad_missing[:8]} "
                f"unexpected: {result.unexpected_keys[:8]}"
            )
        leftover_meta = [
            n for n, p in model.named_parameters() if p.device.type == "meta"
        ]
        if leftover_meta:
            raise RuntimeError(f"Params never loaded: {leftover_meta[:8]}")

        model.generation_config.pad_token_id = 151643
        model.generation_config.eos_token_id = [151645, 151643]
        # built from config, so no sampling defaults were loaded; greedy decode
        # falls into repetition loops on long captions (A-B-A-B forever on
        # low-motion clips). Qwen's recommended sampling for the Qwen3 family:
        model.generation_config.do_sample = True
        # Qwen's recommended sampling: instruct 0.7/0.8, thinking 0.6/0.95
        model.generation_config.temperature = 0.6 if self.is_thinking_model else 0.7
        model.generation_config.top_p = 0.95 if self.is_thinking_model else 0.8
        model.generation_config.top_k = 20
        model.generation_config.repetition_penalty = 1.05

        # swap the slow bf16 Conv3d patch_embed for an equivalent fast linear
        patch_qwen_vl_patch_embed(model)

        if self.caption_config.quantize:
            print(
                "[AITK] Qwen3-Omni loads pre-quantized (convrot8); the quantize "
                "setting is ignored."
            )

        self.model = model
        if self.caption_config.layer_offloading:
            from toolkit.memory_management import MemoryManager

            self.print_and_status_update(
                " - layer offloading enabled: expert banks stay in system RAM, "
                "linears stream per layer"
            )
            # expert banks: stay in system RAM, stream routed experts per call
            for module in model.modules():
                if isinstance(module, ConvRot8Experts):
                    module.enable_offload(self.device_torch)
            # everything the manager doesn't classify must ride to the GPU as
            # unmanaged: the output head, the MoE routers (bare-parameter
            # modules doing F.linear directly), and buffer-only modules
            ignore = [model.lm_head]
            ignore += [
                m
                for m in model.modules()
                if m.__class__.__name__ == "SinusoidsPositionEmbedding"
                or m.__class__.__name__.endswith("TopKRouter")
            ]
            MemoryManager.attach(
                model,
                self.device_torch,
                offload_percent=self.caption_config.layer_offloading_percent,
                ignore_modules=ignore,
            )
        self.model.to(self.device_torch)
        self.processor = AutoProcessor.from_pretrained(self._model_info["base_repo"])
        flush()

    @staticmethod
    def _is_image_file(file_path: str) -> bool:
        return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS

    @staticmethod
    def _is_audio_file(file_path: str) -> bool:
        return os.path.splitext(file_path)[1].lower().lstrip(".") in AUDIO_EXTENSIONS

    def _build_messages(self, _file_path: str):
        if self._is_image_file(_file_path):
            media = {"type": "image", "image": _file_path}
        elif self._is_audio_file(_file_path):
            media = {"type": "audio", "audio": _file_path}
        else:
            media = {"type": "video", "video": _file_path}
        return [
            {
                "role": "user",
                "content": [
                    media,
                    {"type": "text", "text": self.resolve_prompt(self.caption_config.caption_prompt, _file_path)},
                ],
            }
        ]

    def _size_kwargs(self):
        max_pixels = self.caption_config.max_res * self.caption_config.max_res
        # shortest_edge/longest_edge are total pixel counts
        # (min_pixels/max_pixels), not edge lengths
        return {
            "shortest_edge": min(131072, max_pixels),
            "longest_edge": max_pixels,
        }

    def _prep_media(self, file_path: str):
        """CPU side of one file, safe to run in a worker thread: decode +
        subsample frames (or load the image), extract the audio track, render
        the chat text. At batch size 1 the full processor (tokenize, resize,
        mel) runs here too, so the main thread only moves tensors and
        generates."""
        if self._is_image_file(file_path):
            from PIL import Image

            image = Image.open(file_path).convert("RGB")
            item = {"file": file_path, "kind": "image", "image": image, "audio": None}
        elif self._is_audio_file(file_path):
            from transformers.audio_utils import load_audio

            item = {"file": file_path, "kind": "audio", "audio": load_audio(file_path, sampling_rate=16000)}
        else:
            from transformers.video_utils import load_video
            from transformers.audio_utils import load_audio

            frames = load_video(file_path, fps=VIDEO_FPS)
            if isinstance(frames, tuple):
                frames = frames[0]
            audio = None
            try:
                a = load_audio(file_path, sampling_rate=16000)
                if a is not None and a.size > 0:
                    audio = a
            except Exception:
                pass
            item = {
                "file": file_path,
                "kind": "video_audio" if audio is not None else "video_silent",
                "frames": frames,
                "audio": audio,
            }
        template_kwargs = {}
        if self.is_thinking_model and not self.thinking_enabled:
            template_kwargs["enable_thinking"] = False
        item["text"] = self.processor.apply_chat_template(
            self._build_messages(file_path),
            tokenize=False,
            add_generation_prompt=True,
            **template_kwargs,
        )
        if self.caption_config.batch_size <= 1:
            item["inputs"] = self._process_items([item])
        return item

    def _process_items(self, items):
        kind = items[0]["kind"]
        if kind == "image":
            return self.processor(
                text=[it["text"] for it in items],
                images=[it["image"] for it in items],
                return_tensors="pt",
                padding=True,
                size=self._size_kwargs(),
            )
        if kind == "audio":
            return self.processor(
                text=[it["text"] for it in items],
                audio=[it["audio"] for it in items],
                return_tensors="pt",
                padding=True,
            )
        use_audio = kind == "video_audio"
        return self.processor(
            text=[it["text"] for it in items],
            audio=[it["audio"] for it in items] if use_audio else None,
            videos=[it["frames"] for it in items],
            return_tensors="pt",
            padding=True,
            use_audio_in_video=use_audio,
            fps=VIDEO_FPS,
            do_sample_frames=False,
            size=self._size_kwargs(),
        )

    def _caption_batch(self, items):
        """Batched generate over preprocessed items (all the same kind: image,
        video with audio, or silent video). Returns captions in item order."""
        use_audio = items[0]["kind"] == "video_audio"
        if len(items) == 1 and "inputs" in items[0]:
            inputs = items[0]["inputs"]
        else:
            inputs = self._process_items(items)
        inputs = inputs.to(self.device_torch).to(self.torch_dtype)
        # a generate that dies between static-cache creation and its first
        # forward leaves model._cache with uninitialized layers; transformers
        # then raises AttributeError reading cache.max_batch_size on every
        # later call, masking the original error — drop the stale cache
        stale_cache = getattr(self.model, "_cache", None)
        if stale_cache is not None and not stale_cache.is_initialized:
            del self.model._cache
        # under static cache, generate hands the forward a prepared 4D mask;
        # the true 2D padding mask is needed for the prefill rope index
        self.model._pad_mask_2d = inputs.get("attention_mask", None)
        gen_kwargs = self._gen_kwargs(inputs["input_ids"].shape[1])
        if items[0]["kind"] == "audio":
            # greedy lyric transcription loops on musical phrases; stronger penalty to mitigate
            gen_kwargs["repetition_penalty"] = 1.20
        generated_ids = self.model.generate(
            **inputs,
            use_audio_in_video=use_audio,
            **gen_kwargs,
        )
        trimmed = generated_ids[:, inputs["input_ids"].shape[1] :]
        captions = self.processor.batch_decode(
            trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )
        # thinking models emit reasoning first; keep only what follows it
        captions = [c.split("</think>")[-1] if "</think>" in c else c for c in captions]
        return [c.strip() for c in captions]

    def _gen_kwargs(self, input_len: int) -> dict:
        """Generation length controls. Thinking models get their reasoning
        budget on top: max_new_tokens starts counting after </think> closes.
        Under compiled decode, max_length stays constant (fixed cache shape)
        and the real budget lives in the stopping criteria."""
        from transformers.generation import MaxLengthCriteria, StoppingCriteriaList

        max_new = self.caption_config.max_new_tokens
        compiled = self.model.generation_config.cache_implementation == "static"
        criteria = []
        if self.thinking_enabled:
            think_end_id = self.processor.tokenizer.convert_tokens_to_ids("</think>")
            if think_end_id is not None:
                criteria.append(BatchThinkingBudgetCriteria(think_end_id, max_new))
            budget = MAX_THINKING_TOKENS + max_new
        else:
            budget = max_new
        if compiled:
            criteria.append(MaxLengthCriteria(max_length=input_len + budget))
            return {
                "max_length": STATIC_MAX_LENGTH,
                "stopping_criteria": StoppingCriteriaList(criteria),
            }
        kwargs = {"max_new_tokens": budget}
        if criteria:
            kwargs["stopping_criteria"] = StoppingCriteriaList(criteria)
        return kwargs

    def run_caption_loop(self):
        """Batched pipeline: CPU worker threads decode/preprocess videos ahead
        of the GPU, videos are grouped (with-audio vs silent) into batches, and
        each batch runs one model.generate call so decode work is wide enough
        to saturate the GPU."""
        import concurrent.futures
        from collections import deque

        import tqdm as tqdm_mod

        batch_size = max(1, int(self.caption_config.batch_size))
        # smoothing near 1 weights recent files heavily, so the rate estimate
        # recovers quickly after the slow compile-warmup videos
        pbar = tqdm_mod.tqdm(
            total=len(self.file_paths),
            desc="Captioning files",
            unit="file",
            smoothing=0.9,
        )

        def finish(file_path, caption):
            if caption is not None:
                self.save_caption_for_file(file_path, caption)
            self.step_num += 1
            self.update_step()
            pbar.update(1)

        def flush(bucket):
            if len(bucket) == 0:
                return
            items = list(bucket)
            bucket.clear()
            n_real = len(items)
            # keep the batch shape constant for the compiled decode graph:
            # pad a final partial bucket by repeating the last video
            if (
                self.model.generation_config.cache_implementation == "static"
                and 1 < n_real < batch_size
            ):
                items = items + [items[-1]] * (batch_size - n_real)
            try:
                captions = self._caption_batch(items)[:n_real]
                for it, cap in zip(items[:n_real], captions):
                    finish(it["file"], cap)
            except Exception as e:
                print(f"Batch failed ({e}); retrying files individually")
                traceback.print_exc()
                for it in items[:n_real]:
                    finish(it["file"], self.get_caption_for_file(it["file"]))

        executor = concurrent.futures.ThreadPoolExecutor(
            max_workers=max(1, int(self.caption_config.num_workers))
        )
        try:
            futures = deque()
            file_iter = iter(self.file_paths)
            # keep a couple of batches of decode work in flight ahead of the GPU
            lookahead = batch_size * 2 + 2
            for _ in range(lookahead):
                path = next(file_iter, None)
                if path is None:
                    break
                futures.append((path, executor.submit(self._prep_media, path)))

            # batches must be homogeneous: the processor call differs per kind
            buckets = {"image": [], "audio": [], "video_audio": [], "video_silent": []}
            while futures:
                if self.is_ui_captioner:
                    self.maybe_stop()
                    if self.is_stopping:
                        break
                path, fut = futures.popleft()
                nxt = next(file_iter, None)
                if nxt is not None:
                    futures.append((nxt, executor.submit(self._prep_media, nxt)))
                try:
                    item = fut.result()
                except Exception as e:
                    print(f"Error preprocessing {path}: {e}")
                    finish(path, None)
                    continue
                bucket = buckets[item["kind"]]
                bucket.append(item)
                if len(bucket) >= batch_size:
                    flush(bucket)
            for bucket in buckets.values():
                flush(bucket)
        finally:
            executor.shutdown(wait=False, cancel_futures=True)
            pbar.close()

    def maybe_compile_models(self):
        """CUDA-graph decode: static kv cache + reduce-overhead compile of the
        text model. Each decode step replays as one captured graph, removing
        the per-kernel python/launch gaps that cap GPU utilization at small
        batch sizes. First video per batch shape is slow (compile warmup)."""
        if not self.caption_config.compile:
            return
        if self.caption_config.layer_offloading:
            # cuda graphs need every tensor GPU-resident; offloaded weights
            # live in system RAM, so the compiled decode path cannot capture
            print("[AITK] layer offloading is on; skipping compiled decode.")
            return
        import importlib.util

        if importlib.util.find_spec("triton") is None:
            print("[AITK] compile requested but triton is not installed, skipping.")
            return
        # a static (compileable) cache makes generate auto-compile its decode
        # loop into one cuda graph; prefill stays eager. Per-block graphs were
        # tried and don't compose (graph capture must own the in-place kv-cache
        # writes, and cudagraph trees can't span 48 independent graphs), and
        # fusion-only block compile doesn't touch the launch gaps that matter.
        # With prepare_inputs_for_generation stripping per-video media shapes
        # from decode steps, this compiles exactly once and caches to disk.
        self.model.generation_config.cache_implementation = "static"
        print(
            "[AITK] Compiled decode enabled (static cache + cuda graphs). "
            "The first video compiles (~2 min cold, faster once cached)."
        )

    def get_caption_for_file(self, file_path: str) -> str:
        # single-file path (and the per-file fallback when a batch fails):
        # same prep + generate flow as the batched loop, for one item
        try:
            return self._caption_batch([self._prep_media(file_path)])[0]
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
            traceback.print_exc()
            return None