text_encoder/text_encoder_mage_flow.py

"""Mage-Flow text encoder (Qwen3-VL, packed varlen conditioning).

Vendored from microsoft/Mage (`mage_flow`, MIT) at commit 76bec2bb3818, with a diffusers-convention
wrapper appended. Upstream is the reference implementation: the numerics here are its own functions,
not a reimplementation. The mandatory content-policy gate upstream runs in ``generate_images`` is not
part of this port.

Copyright (c) 2026 Microsoft. Licensed under the MIT License.
"""

from __future__ import annotations


import os
from collections.abc import Callable
from dataclasses import dataclass

try:
    from typing import Unpack
except ImportError:
    from typing_extensions import Unpack

import torch
from torch import nn
from transformers import AutoProcessor, AutoTokenizer, Cache, Qwen3VLForConditionalGeneration
from transformers.cache_utils import DynamicCache
from transformers.masking_utils import create_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.models.qwen3_vl.modeling_qwen3_vl import (
    Qwen3VLCausalLMOutputWithPast,
    apply_rotary_pos_emb,
    eager_attention_forward,
)
from transformers.utils import ModelOutput



# ===========================================================================
# Custom Qwen3-VL model (customizable forward output)
# ===========================================================================

import logging

logger = logging.getLogger(__name__)

"""Attention backend shim — switchable between Flash Attention 2 and 4.

Exports a single ``flash_attn_varlen_func`` with the FA2 calling convention.
The underlying kernel is selected at runtime via ``set_attn_backend(name)``
(default: ``"flash2"``). The selected kernel is resolved lazily on the first
call so model-config-driven selection (which happens after this module is
imported) takes effect.

Modules that previously did ``from flash_attn import flash_attn_varlen_func``
should import from here instead.

For the FA4 path, calling-convention differences are normalised:

* ``window_size=(-1, -1)`` (FA2 "no window") -> ``(None, None)`` (FA4).
* ``block_table`` -> ``page_table``.
* FA4's optional ``(out, lse)`` tuple return is unwrapped to ``out``.
* ``dropout_p>0`` / ``alibi_slopes`` / ``return_attn_probs`` raise on FA4.
"""


from typing import Any, Callable

_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}

_BACKEND: str = "flash2"
_RESOLVED_FN: Callable[..., Any] | None = None


def _normalize(name: str) -> str:
    n = name.lower().strip()
    if n in _FA2_ALIASES:
        return "flash2"
    if n in _FA4_ALIASES:
        return "flash4"
    if n in _SDPA_ALIASES:
        return "sdpa"
    raise ValueError(
        f"Unknown attention backend {name!r}; expected one of "
        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
    )


def set_attn_backend(name: str) -> None:
    """Select the flash-attn backend used by ``flash_attn_varlen_func``.

    Safe to call multiple times; clears the cached resolution on change.
    """
    global _BACKEND, _RESOLVED_FN
    new = _normalize(name)
    if new != _BACKEND:
        _RESOLVED_FN = None
    _BACKEND = new



def _resolve_fa2() -> Callable[..., Any]:
    # Imported by name, not with a plain ``import``: transformers' dynamic-module loader scans this file
    # and refuses to load it when it sees an import of a package that is not installed — even one inside
    # a function that the sdpa fallback never reaches.
    import importlib

    return importlib.import_module("flash_attn").flash_attn_varlen_func


def _resolve_fa4() -> Callable[..., Any]:
    # 같은 이유로 이름으로 import (정적 스캔이 하드 요구로 보지 않게)
    import importlib

    _fa4_fn = importlib.import_module("flash_attn.cute").flash_attn_varlen_func

    def _fa4_wrapper(
        q,
        k,
        v,
        cu_seqlens_q=None,
        cu_seqlens_k=None,
        max_seqlen_q=None,
        max_seqlen_k=None,
        dropout_p: float = 0.0,
        softmax_scale=None,
        causal: bool = False,
        window_size=(-1, -1),
        softcap: float = 0.0,
        alibi_slopes=None,
        deterministic: bool = False,
        return_attn_probs: bool = False,
        block_table=None,
        **_unused: Any,
    ):
        if dropout_p and dropout_p > 0:
            raise NotImplementedError("FA4 backend does not support dropout_p>0")
        if alibi_slopes is not None:
            raise NotImplementedError("FA4 backend does not support alibi_slopes")
        if return_attn_probs:
            raise NotImplementedError("FA4 backend does not support return_attn_probs")

        win_l, win_r = window_size
        if win_l == -1:
            win_l = None
        if win_r == -1:
            win_r = None

        out = _fa4_fn(
            q,
            k,
            v,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=max_seqlen_k,
            softmax_scale=softmax_scale,
            causal=causal,
            window_size=(win_l, win_r),
            softcap=softcap,
            deterministic=deterministic,
            page_table=block_table,
            return_lse=False,
        )
        if isinstance(out, tuple):
            out = out[0]
        return out

    return _fa4_wrapper


def _resolve_sdpa() -> Callable[..., Any]:
    """FA2 varlen → per-sequence torch.SDPA fallback.

    Use when flash-attn is unavailable (e.g. CUDA 13 has no prebuilt wheel
    and source build is brittle). Slower than FA2 (one SDPA dispatch per
    sequence), but functionally equivalent for the dense / causal / no-alibi
    paths mageflow actually uses. Window / softcap / alibi / paged-attn /
    return_attn_probs are not supported and will raise.
    """
    import torch
    import torch.nn.functional as F

    def _sdpa_wrapper(
        q,
        k,
        v,
        cu_seqlens_q=None,
        cu_seqlens_k=None,
        max_seqlen_q=None,
        max_seqlen_k=None,
        dropout_p: float = 0.0,
        softmax_scale=None,
        causal: bool = False,
        window_size=(-1, -1),
        softcap: float = 0.0,
        alibi_slopes=None,
        deterministic: bool = False,
        return_attn_probs: bool = False,
        block_table=None,
        **_unused: Any,
    ):
        if dropout_p and dropout_p > 0:
            raise NotImplementedError("SDPA backend does not support dropout_p>0")
        if alibi_slopes is not None:
            raise NotImplementedError("SDPA backend does not support alibi_slopes")
        if return_attn_probs:
            raise NotImplementedError("SDPA backend does not support return_attn_probs")
        if softcap and softcap > 0:
            raise NotImplementedError("SDPA backend does not support softcap")
        if window_size not in ((-1, -1), (None, None), (0, 0)):
            raise NotImplementedError(
                f"SDPA backend does not support sliding window (got {window_size})"
            )
        if block_table is not None:
            raise NotImplementedError("SDPA backend does not support paged attention")
        if cu_seqlens_q is None or cu_seqlens_k is None:
            raise ValueError("SDPA backend requires cu_seqlens_q and cu_seqlens_k")

        # GQA: FA2 broadcasts k/v across query head groups natively; torch SDPA
        # does not (the q vs k head-dim mismatch is the AssertionError "tensor
        # a (32) must match tensor b (8) at non-singleton dimension 1" we'd see
        # otherwise). Repeat k/v along the head dim to match q before the loop.
        n_heads_q = q.shape[1]
        n_heads_kv = k.shape[1]
        if n_heads_q != n_heads_kv:
            if n_heads_q % n_heads_kv != 0:
                raise ValueError(
                    f"SDPA backend GQA expansion requires q heads ({n_heads_q}) "
                    f"to be divisible by k/v heads ({n_heads_kv})"
                )
            repeat = n_heads_q // n_heads_kv
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)

        # q/k/v: (total_tokens, nheads, head_dim). Dispatch SDPA per sequence,
        # then concat. Python-level loop is fine since nseq is small (one per
        # image in the pack) and image-gen latency is dominated by sampling.
        cu_q = cu_seqlens_q.tolist()
        cu_k = cu_seqlens_k.tolist()
        outs = []
        for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]):
            # (s, h, d) → (1, h, s, d)
            q_i = q[qs:qe].transpose(0, 1).unsqueeze(0)
            k_i = k[ks:ke].transpose(0, 1).unsqueeze(0)
            v_i = v[ks:ke].transpose(0, 1).unsqueeze(0)
            out_i = F.scaled_dot_product_attention(
                q_i,
                k_i,
                v_i,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=causal,
                scale=softmax_scale,
            )
            # (1, h, s, d) → (s, h, d)
            outs.append(out_i.squeeze(0).transpose(0, 1))
        return torch.cat(outs, dim=0).contiguous()

    return _sdpa_wrapper


def _resolve() -> Callable[..., Any]:
    global _RESOLVED_FN
    if _RESOLVED_FN is None:
        if _BACKEND == "flash4":
            _RESOLVED_FN = _resolve_fa4()
        elif _BACKEND == "sdpa":
            _RESOLVED_FN = _resolve_sdpa()
        else:
            try:
                _RESOLVED_FN = _resolve_fa2()
            except ImportError:
                # flash-attn 2 needs sm80+ and a matching build; sdpa is the portable varlen path, so a
                # missing kernel falls back instead of failing the load.
                logger.warning("flash-attn 2 is unavailable; using the sdpa attention backend")
                _RESOLVED_FN = _resolve_sdpa()
    return _RESOLVED_FN


def flash_attn_varlen_func(*args, **kwargs):
    return _resolve()(*args, **kwargs)


__all__ = ["flash_attn_varlen_func", "set_attn_backend"]


@dataclass
class Qwen3VLModelOutput(ModelOutput):
    """Flexible output class for custom Qwen3-VL model."""

    loss: torch.FloatTensor | None = None
    logits: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    last_hidden_state: torch.FloatTensor | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None
    rope_deltas: torch.LongTensor | None = None


class CustomQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
    """
    Custom Qwen3-VL model that allows customizing the forward output.

    This class inherits from Qwen3VLForConditionalGeneration and provides
    hooks to customize what is returned from the forward pass.

    Example usage:
        ```python
        model = CustomQwen3VLForConditionalGeneration.from_pretrained(
            "Qwen/Qwen3-VL-8B-Instruct",
            attn_implementation="flash_attention_2"  # Use flash attention for faster inference
        )

        # Option 1: Use built-in output modes
        model.set_output_mode("embedding")  # Only return last hidden state (default)
        model.set_output_mode("full")       # Return everything
        model.set_output_mode("logits")     # Only return logits

        # Option 2: Set a custom output processor
        def my_custom_output(hidden_states, logits, outputs, **kwargs):
            return {"embeddings": hidden_states, "pooled": hidden_states.mean(dim=1)}
        model.set_output_processor(my_custom_output)
        ```
    """

    # Output mode constants
    OUTPUT_MODE_FULL = "full"
    OUTPUT_MODE_EMBEDDING = "embedding"
    OUTPUT_MODE_LOGITS = "logits"
    OUTPUT_MODE_HIDDEN = "hidden"

    def __init__(self, config):
        super().__init__(config)
        self._output_mode = self.OUTPUT_MODE_EMBEDDING
        self._skip_lm_head = True

    def set_output_mode(self, mode: str):
        """
        Set the output mode for the forward pass.

        Args:
            mode: One of:
                - "full": Return full Qwen3VLCausalLMOutputWithPast
                - "embedding": Only return last hidden state (skip lm_head) (default)
                - "logits": Only return logits
                - "hidden": Return all hidden states
        """
        valid_modes = [
            self.OUTPUT_MODE_FULL,
            self.OUTPUT_MODE_EMBEDDING,
            self.OUTPUT_MODE_LOGITS,
            self.OUTPUT_MODE_HIDDEN,
        ]
        if mode not in valid_modes:
            raise ValueError(f"Invalid output mode: {mode}. Must be one of {valid_modes}")
        self._output_mode = mode
        self._skip_lm_head = mode == self.OUTPUT_MODE_EMBEDDING

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        pixel_values: torch.Tensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ) -> Qwen3VLCausalLMOutputWithPast | Qwen3VLModelOutput | dict | torch.Tensor:
        """
        Forward pass with customizable output.

        Returns different outputs based on the configured output mode or custom processor.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )

        # Get outputs from the base model (Qwen3VLModel)
        outputs = self.model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )

        # Get the last hidden state
        hidden_states = outputs[0]  # This is the last hidden state

        # Compute logits if not skipping lm_head
        logits = None
        if not self._skip_lm_head:
            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
            logits = self.lm_head(hidden_states[:, slice_indices, :])

        # Compute loss if labels are provided
        loss = None
        if labels is not None and logits is not None:
            loss = self.loss_function(
                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
            )

        # Return based on output mode
        if self._output_mode == self.OUTPUT_MODE_EMBEDDING:
            return Qwen3VLModelOutput(
                last_hidden_state=hidden_states,
                past_key_values=outputs.past_key_values,
                attentions=outputs.attentions,
                rope_deltas=outputs.rope_deltas,
            )
        elif self._output_mode == self.OUTPUT_MODE_LOGITS:
            return logits
        elif self._output_mode == self.OUTPUT_MODE_HIDDEN:
            return Qwen3VLModelOutput(
                last_hidden_state=hidden_states,
                hidden_states=outputs.hidden_states,
                past_key_values=outputs.past_key_values,
                attentions=outputs.attentions,
                rope_deltas=outputs.rope_deltas,
            )
        else:  # OUTPUT_MODE_FULL
            return Qwen3VLCausalLMOutputWithPast(
                loss=loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
                rope_deltas=outputs.rope_deltas,
            )


# ===========================================================================
# Packing-aware forward patches (cu_seqlens) for the Qwen3-VL text encoder
# ===========================================================================

def model_forward(
    self,
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    use_cache: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    # args for deepstack
    visual_pos_masks: torch.Tensor | None = None,
    deepstack_visual_embeds: list[torch.Tensor] | None = None,
    **kwargs: Unpack[FlashAttentionKwargs],
) -> tuple | BaseModelOutputWithPast:
    r"""
    visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*):
        The mask of the visual positions.
    deepstack_visual_embeds (`list[torch.Tensor]`, *optional*):
        The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim).
        The feature is extracted from the different visual encoder layers, and fed to the decoder
        hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334).
    """
    if (input_ids is None) ^ (inputs_embeds is not None):
        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

    # torch.jit.trace() doesn't support cache objects in the output
    if use_cache and past_key_values is None and not torch.jit.is_tracing():
        past_key_values = DynamicCache(config=self.config)

    if inputs_embeds is None:
        inputs_embeds = self.embed_tokens(input_ids)

    if cache_position is None:
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
        cache_position = torch.arange(
            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
        )

    # the hard coded `3` is for temporal, height and width.
    if position_ids is None:
        position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
    elif position_ids.ndim == 2:
        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)

    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
        text_position_ids = position_ids[0]
        position_ids = position_ids[1:]
    else:
        text_position_ids = position_ids[0]

    if kwargs.get("cu_seqlens") is None:
        attention_mask = create_causal_mask(
            config=self.config,
            input_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            past_key_values=past_key_values,
            position_ids=text_position_ids,
        )

    hidden_states = inputs_embeds

    # create position embeddings to be shared across the decoder layers
    position_embeddings = self.rotary_emb(hidden_states, position_ids)

    # decoder layers
    for layer_idx, decoder_layer in enumerate(self.layers):
        layer_outputs = decoder_layer(
            hidden_states,
            attention_mask=attention_mask,
            position_ids=text_position_ids,
            past_key_values=past_key_values,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = layer_outputs

        # add visual features to the hidden states of first several layers
        if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
            hidden_states = self._deepstack_process(
                hidden_states,
                visual_pos_masks,
                deepstack_visual_embeds[layer_idx],
            )

    hidden_states = self.norm(hidden_states)

    return BaseModelOutputWithPast(
        last_hidden_state=hidden_states,
        past_key_values=past_key_values,
    )


def forward(
    self,
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    past_key_values: Cache | None = None,
    cache_position: torch.LongTensor | None = None,
    **kwargs: Unpack[FlashAttentionKwargs],
) -> tuple[torch.Tensor, torch.Tensor | None]:
    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)

    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    if past_key_values is not None:
        # sin and cos are specific to RoPE models; cache_position needed for the static cache
        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

    cu_seqlens = kwargs.get("cu_seqlens", None)

    if cu_seqlens is None:
        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )
    else:
        max_seqlen = torch.diff(cu_seqlens).max().item() if cu_seqlens is not None else None
        query_states = query_states.transpose(1, 2).squeeze(0)
        key_states = key_states.transpose(1, 2).squeeze(0)
        value_states = value_states.transpose(1, 2).squeeze(0)
        attn_output = flash_attn_varlen_func(
            q=query_states,
            k=key_states,
            v=value_states,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_k=cu_seqlens,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            causal=True,
            window_size=(-1, -1),
            softmax_scale=self.head_dim**-0.5,
            dropout_p=0.0,
        )

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, None


def qwen3_patch_forward():
    """Patch the Qwen3-VL text model + attention forwards to support packed
    varlen (cu_seqlens) inputs used by ``TextEncoder.forward``."""
    from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention, Qwen3VLTextModel

    Qwen3VLTextModel.forward = model_forward
    Qwen3VLTextAttention.forward = forward


# ===========================================================================
# TextEncoder wrapper (packed text -> DiT conditioning embeddings)
# ===========================================================================
_FA2_ALIASES = {"flash2", "fa2", "flash_attention_2", "flash_attn_2"}
_FA4_ALIASES = {"flash4", "fa4", "flash_attention_4", "flash_attn_4"}
_SDPA_ALIASES = {"sdpa", "torch_sdpa", "scaled_dot_product_attention"}


def _resolve_hf_attn_impl(attn_type: str) -> str:
    """Map a project-level attn_type to a HuggingFace ``attn_implementation`` string.

    ``VF_HF_ATTN_IMPL`` env var, if set, takes precedence (useful for forcing
    sdpa on machines without flash-attn). For FA4 we additionally probe that
    the CUTE-DSL kernel is importable and (when available) ask the HF helper
    to confirm; if not, fall back to sdpa rather than crashing at load time.
    """
    override = os.environ.get("VF_HF_ATTN_IMPL")
    if override:
        return override

    name = attn_type.lower().strip()
    if name in _FA2_ALIASES:
        return "flash_attention_2"
    if name in _FA4_ALIASES:
        try:
            import flash_attn.cute  # noqa: F401
            fa4_importable = True
        except Exception:
            fa4_importable = False
        if fa4_importable:
            try:
                from transformers.utils.import_utils import is_flash_attn_4_available
                if is_flash_attn_4_available():
                    return "flash_attention_4"
            except ImportError:
                return "flash_attention_4"
        logger.warning(
            "attn_type=flash4 requested but flash_attn.cute is unavailable; "
            "falling back to sdpa for HF text encoder."
        )
        return "sdpa"
    if name in _SDPA_ALIASES:
        return "sdpa"
    raise ValueError(
        f"Unknown attn_type {attn_type!r}; expected one of "
        f"{sorted(_FA2_ALIASES | _FA4_ALIASES | _SDPA_ALIASES)}"
    )


SEQ_MULTI_OF = 32





# ---------------------------------------------------------------------------
# transformers-version shim + diffusers component wrapper
# ---------------------------------------------------------------------------
# ``create_causal_mask`` renamed ``input_embeds`` to ``inputs_embeds`` and dropped ``cache_position``
# after the transformers release this code was written against, so the call above is translated here
# rather than edited upstream.
_upstream_create_causal_mask = create_causal_mask


def _create_causal_mask(*args, **kwargs):
    if "input_embeds" in kwargs:
        kwargs["inputs_embeds"] = kwargs.pop("input_embeds")
    if "cache_position" in kwargs and "cache_position" not in _CREATE_CAUSAL_MASK_PARAMS:
        kwargs.pop("cache_position")
    return _upstream_create_causal_mask(*args, **kwargs)


import inspect  # noqa: E402

_CREATE_CAUSAL_MASK_PARAMS = set(inspect.signature(_upstream_create_causal_mask).parameters)
create_causal_mask = _create_causal_mask

qwen3_patch_forward()


class MageFlowTextEncoder(CustomQwen3VLForConditionalGeneration):
    """Qwen3-VL text encoder with Mage-Flow's packed (varlen) conditioning forward.

    ``encode_packed`` is upstream's ``TextEncoder.forward`` body: several prompts are concatenated and
    isolated by ``cu_seqlens`` in one launch, each sequence's leading template tokens are dropped, and
    the pooled vector is the mean over what remains.
    """

    def encode_packed(self, input_ids, cu_seqlens, drop_idx: int = 0, inputs: dict | None = None) -> dict:
        seqlens_list = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().tolist()
        position_ids = torch.cat([torch.arange(_length, device=input_ids.device) for _length in seqlens_list])

        forward_kwargs = {
            "input_ids": input_ids.unsqueeze(0).to(self.device),
            "cu_seqlens": cu_seqlens,
            "position_ids": position_ids.unsqueeze(0).to(self.device),
            "output_hidden_states": False,
            "max_seqlen": None,
        }
        if inputs is not None:
            for _key in ("pixel_values", "image_grid_thw"):
                if inputs.get(_key, None) is not None:
                    forward_kwargs[_key] = inputs[_key].to(self.device)

        with torch.no_grad():
            outputs = self(**forward_kwargs)
        hidden = outputs.last_hidden_state if getattr(outputs, "last_hidden_state", None) is not None \
            else outputs.hidden_states[-1]
        hidden = hidden.squeeze(0)

        txt_list, vec_list, valid_lengths = list(), list(), list()
        for _hidden in torch.split(hidden, seqlens_list, dim=0):
            _valid = _hidden[drop_idx:]
            txt_list.append(_valid)
            vec_list.append(_valid.mean(dim=0))
            valid_lengths.append(_valid.shape[0])
        return {
            "txt": torch.cat(txt_list, dim=0),
            "vec": torch.stack(vec_list, dim=0),
            "txt_seq_lens": torch.tensor(valid_lengths, device=input_ids.device),
        }