modeling_mineru_diffusion.py

import re, time
from typing import Optional, Union

import torch
import torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VisionTransformerPretrainedModel

from .configuration_mineru_diffusion import MinerUDiffusionConfig, SDARConfig

try:
    from flash_attn import flash_attn_func
except ImportError:
    flash_attn_func = None


def _new_dynamic_cache(config: Optional[SDARConfig] = None) -> DynamicCache:
    try:
        if config is not None:
            return DynamicCache(config=config)
    except TypeError:
        pass
    return DynamicCache()


class PatchMerger(nn.Module):
    def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
        super().__init__()
        self.hidden_size = context_dim * (spatial_merge_size**2)
        self.ln_q = nn.LayerNorm(context_dim, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.GELU(),
            nn.Linear(self.hidden_size, dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.ln_q(x).view(-1, self.hidden_size)
        return self.mlp(x)


def build_projection(projection_type: str, in_dim: int, out_dim: int) -> nn.Module:
    pm_match = re.match(r"(?:patch_merger|pm)(\d+)x$", projection_type)
    if pm_match:
        merge_size = int(pm_match.group(1))
        return PatchMerger(out_dim, in_dim, spatial_merge_size=merge_size)

    raise ValueError(f"Only patch_merger-style projectors are supported, got: {projection_type}")


class PerceiverProjection(nn.Module):
    def __init__(self, projection_type: str, in_dim: int, out_dim: int):
        super().__init__()
        self.projection = build_projection(projection_type, in_dim, out_dim)

    def forward(self, input_embeds: torch.Tensor) -> torch.Tensor:
        return self.projection(input_embeds)


class SDARRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class SDARMLP(nn.Module):
    def __init__(self, config: SDARConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


class SDARAttention(nn.Module):
    def __init__(self, config: SDARConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.scaling = self.head_dim**-0.5
        self.is_causal = False
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
        self.q_norm = SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def _project_hidden_states(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        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)
        return query_states, key_states, value_states

    def _use_initialized_cache(self, past_key_value: Cache) -> bool:
        cache_layers = getattr(past_key_value, "layers", None)
        return (
            cache_layers is not None
            and len(cache_layers) > self.layer_idx
            and cache_layers[self.layer_idx].is_initialized
        )

    def _update_past_key_values(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        past_key_value: Optional[Cache],
        store_kv: bool,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if past_key_value is None:
            return key_states, value_states

        if store_kv:
            return past_key_value.update(key_states, value_states, self.layer_idx)

        if self._use_initialized_cache(past_key_value):
            cache_layer = past_key_value.layers[self.layer_idx]
            past_key_states = cache_layer.keys
            past_value_states = cache_layer.values
        elif len(past_key_value) > self.layer_idx:
            past_key_states, past_value_states = past_key_value[self.layer_idx]
        else:
            return key_states, value_states

        key_states = torch.cat([past_key_states, key_states], dim=-2)
        value_states = torch.cat([past_value_states, value_states], dim=-2)
        return key_states, value_states

    def _can_use_flash_attention(
        self,
        query_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> bool:
        return (
            flash_attn_func is not None
            and query_states.device.type == "cuda"
            and attention_mask is not None
            and torch.all(attention_mask)
            and query_states.dtype in (torch.float16, torch.bfloat16)
        )

    def _flash_attention_forward(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        input_shape: torch.Size,
    ) -> torch.Tensor:
        attn_output = flash_attn_func(
            query_states.transpose(1, 2),
            key_states.transpose(1, 2),
            value_states.transpose(1, 2),
            causal=self.is_causal,
            softmax_scale=self.scaling,
        )
        return attn_output.reshape(*input_shape, -1).contiguous()

    def _sdpa_attention_forward(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
        input_shape: torch.Size,
    ) -> torch.Tensor:
        attn_output = F.scaled_dot_product_attention(
            query=query_states,
            key=key_states,
            value=value_states,
            attn_mask=attention_mask.bool() if attention_mask is not None else None,
            is_causal=self.is_causal,
            scale=self.scaling,
            enable_gqa=True,
        )
        return attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings,
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        position_ids: Optional[torch.LongTensor] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        store_kv: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        input_shape = hidden_states.shape[:-1]
        query_states, key_states, value_states = self._project_hidden_states(hidden_states)
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
        key_states, value_states = self._update_past_key_values(key_states, value_states, past_key_value, store_kv)

        if self._can_use_flash_attention(query_states, attention_mask):
            attn_output = self._flash_attention_forward(query_states, key_states, value_states, input_shape)
        else:
            attn_output = self._sdpa_attention_forward(query_states, key_states, value_states, attention_mask, input_shape)
        return self.o_proj(attn_output), None


class SDARDecoderLayer(nn.Module):
    def __init__(self, config: SDARConfig, layer_idx: int):
        super().__init__()
        self.self_attn = SDARAttention(config=config, layer_idx=layer_idx)
        self.mlp = SDARMLP(config)
        self.input_layernorm = SDARRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = SDARRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        position_ids=None,
        past_key_value=None,
        output_attentions: bool = False,
        use_cache: bool = False,
        store_kv: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings=None,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            position_embeddings=position_embeddings,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            store_kv=store_kv,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = residual + hidden_states
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        outputs = (hidden_states,)
        if output_attentions:
            outputs += (self_attn_weights,)
        return outputs


class SDARRotaryEmbedding(nn.Module):
    def __init__(self, config: SDARConfig, device=None):
        super().__init__()
        self.config = config
        if config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
        else:
            self.rope_type = "default"
        if self.rope_type == "default" or self.rope_type not in ROPE_INIT_FUNCTIONS:
            dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
            base = float(getattr(config, "rope_theta", 10000.0))
            inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
            self.attention_scaling = 1.0
        else:
            inv_freq, self.attention_scaling = ROPE_INIT_FUNCTIONS[self.rope_type](config, device)
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @dynamic_rope_update
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
        position_ids_expanded = position_ids[:, None, :].float()
        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with torch.autocast(device_type=device_type, enabled=False):
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


class SDARPreTrainedModel(PreTrainedModel):
    config_class = SDARConfig
    base_model_prefix = "model"
    _no_split_modules = ["SDARDecoderLayer"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, SDARRMSNorm):
            module.weight.data.fill_(1.0)


class SDARModel(SDARPreTrainedModel):
    def __init__(self, config: SDARConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([SDARDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm = SDARRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = SDARRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        store_kv: bool = False,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = True,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        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
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if not isinstance(past_key_values, (type(None), Cache)):
            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")

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

        if use_cache and past_key_values is None:
            past_key_values = _new_dynamic_cache(self.config)
        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,
            )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        hidden_states = inputs_embeds
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            layer_outputs = layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                position_embeddings=position_embeddings,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                store_kv=store_kv,
                cache_position=cache_position,
                **kwargs,
            )
            hidden_states = layer_outputs[0]
            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )


class SDARForCausalLM(SDARPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: SDARConfig):
        super().__init__(config)
        self.model = SDARModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @staticmethod
    def _build_block_attention_mask(num_blocks: int, block_length: int, device: torch.device) -> torch.Tensor:
        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device, dtype=torch.bool))
        return block_mask.repeat_interleave(block_length, dim=0).repeat_interleave(block_length, dim=1)

    def _build_full_attention_mask(
        self,
        prompt_length: int,
        gen_length: int,
        block_length: int,
        device: torch.device,
    ) -> torch.Tensor:
        prompt_blocks = (prompt_length + block_length - 1) // block_length
        prompt_mask = self._build_block_attention_mask(prompt_blocks, block_length, device)
        prompt_mask = prompt_mask[-prompt_length:, -prompt_length:]

        gen_blocks = gen_length // block_length
        gen_mask = self._build_block_attention_mask(gen_blocks, block_length, device)

        full_attn_mask = torch.zeros(
            prompt_length + gen_length,
            prompt_length + gen_length,
            device=device,
            dtype=torch.bool,
        )
        full_attn_mask[:prompt_length, :prompt_length] = prompt_mask
        full_attn_mask[prompt_length:, :prompt_length] = True
        full_attn_mask[prompt_length:, prompt_length:] = gen_mask
        return full_attn_mask

    def _initialize_generation_buffers(
        self,
        inputs_embeds: torch.Tensor,
        gen_length: int,
        mask_token_id: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        batch_size = inputs_embeds.size(0)
        prompt_length = inputs_embeds.size(1)
        device = inputs_embeds.device
        mask_token = torch.tensor([mask_token_id], device=device)
        mask_embeds = self.get_input_embeddings()(mask_token)
        masked_generation_embeds = mask_embeds.unsqueeze(0).expand(batch_size, gen_length, -1)
        x_embeds = torch.cat([inputs_embeds, masked_generation_embeds], dim=1)
        tokens = torch.full(
            (batch_size, prompt_length + gen_length),
            mask_token_id,
            dtype=torch.long,
            device=device,
        )
        step_map = torch.zeros_like(tokens, dtype=torch.int64)
        step_time = torch.zeros_like(tokens, dtype=torch.float)
        return x_embeds, tokens, step_map, step_time, mask_embeds

    @staticmethod
    def _prepare_stop_tokens(
        stopping_criteria,
        tokenizer,
        device: torch.device,
    ) -> list[torch.Tensor]:
        if stopping_criteria is None:
            return []
        if tokenizer is None:
            raise ValueError("tokenizer is required when stopping_criteria is not None")
        return [
            torch.tensor(tokenizer.encode(stop_str, add_special_tokens=False), device=device)
            for stop_str in stopping_criteria
        ]

    def _select_transfer_index(
        self,
        confidence: torch.Tensor,
        threshold: float,
        transfer_count: int,
    ) -> torch.Tensor:
        transfer_index = torch.zeros_like(confidence, dtype=torch.bool)
        for batch_idx in range(confidence.shape[0]):
            high_confidence = confidence[batch_idx] > threshold
            if high_confidence.sum() >= transfer_count:
                transfer_index[batch_idx] = high_confidence
                continue
            _, top_indices = torch.topk(confidence[batch_idx], transfer_count)
            transfer_index[batch_idx, top_indices] = True
        return transfer_index

    @staticmethod
    def _find_stop_position(
        generated_tokens: torch.Tensor,
        stop_tokens: list[torch.Tensor],
    ) -> Optional[int]:
        for stop_token in stop_tokens:
            stop_length = stop_token.numel()
            if stop_length == 0 or generated_tokens.numel() < stop_length:
                continue
            for end_idx in range(stop_length, generated_tokens.size(0) + 1):
                if torch.equal(generated_tokens[end_idx - stop_length : end_idx], stop_token):
                    return end_idx - stop_length
        return None

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = True,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            return_dict=True,
            **kwargs,
        )
        hidden_states = outputs.last_hidden_state
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        hidden_states = hidden_states[:, slice_indices, :].contiguous()
        logits = self.lm_head(hidden_states)
        if not return_dict:
            return (logits, outputs.past_key_values)
        return CausalLMOutputWithPast(
            loss=None,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=None,
            attentions=None,
        )

    @staticmethod
    def top_k_logits(logits, k):
        if k <= 0:
            return logits
        values, _ = torch.topk(logits, k)
        min_values = values[..., -1, None]
        return torch.where(logits < min_values, torch.full_like(logits, float("-inf")), logits)

    @staticmethod
    def top_p_logits(logits, p):
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_mask = cumulative_probs > p
        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
        sorted_mask[..., 0] = False
        mask_indices = torch.scatter(torch.full_like(logits, False, dtype=torch.bool), -1, sorted_indices, sorted_mask)
        return logits.masked_fill(mask_indices, float("-inf"))

    def sample_with_temperature_topk_topp(self, logits, temperature=1.0, top_k=0, top_p=1.0):
        orig_shape = logits.shape[:-1]
        vocab_size = logits.shape[-1]
        logits = logits.reshape(-1, vocab_size)
        if temperature != 1.0:
            logits = logits / temperature
        if top_k > 0:
            logits = self.top_k_logits(logits, top_k)
        if top_p < 1.0:
            logits = self.top_p_logits(logits, top_p)
        probs = F.softmax(logits, dim=-1)
        token = torch.multinomial(probs, num_samples=1)
        token_prob = torch.gather(probs, -1, token)
        return token.view(*orig_shape), token_prob.view(*orig_shape)

    @staticmethod
    def get_num_transfer_tokens(block_length, steps):
        base = block_length // steps
        remainder = block_length % steps
        num_transfer_tokens = torch.zeros(steps, dtype=torch.int64) + base
        num_transfer_tokens[:remainder] += 1
        return num_transfer_tokens

    def generate_with_embeds(
        self,
        inputs_embeds,
        gen_length,
        block_length,
        mask_token_id,
        denoising_steps=8,
        temperature=1.0,
        top_k=0,
        top_p=1.0,
        remasking_strategy="low_confidence_dynamic",
        dynamic_threshold=0.85,
        stopping_criteria=None,
        tokenizer=None,
        **kwargs,
    ):
        if gen_length % block_length != 0:
            raise ValueError(f"gen_length({gen_length}) must be multiple of block_length({block_length})")
        if remasking_strategy != "low_confidence_dynamic":
            raise ValueError("Only remasking_strategy='low_confidence_dynamic' is supported.")

        prompt_length = inputs_embeds.size(1)
        past_key_values = _new_dynamic_cache(self.config)
        gen_blocks = gen_length // block_length
        full_attn_mask = self._build_full_attention_mask(prompt_length, gen_length, block_length, inputs_embeds.device)
        position_ids = torch.arange(0, prompt_length + gen_length, device=inputs_embeds.device).unsqueeze(0)
        x_embeds, x, step_map, step_time, mask_embeds = self._initialize_generation_buffers(inputs_embeds, gen_length, mask_token_id)

        if prompt_length > 0:
            prompt_attn_mask = full_attn_mask[:prompt_length, :prompt_length].unsqueeze(0).unsqueeze(0)
            self(
                inputs_embeds=x_embeds[:, :prompt_length, :],
                attention_mask=prompt_attn_mask,
                position_ids=position_ids[:, :prompt_length],
                past_key_values=past_key_values,
                use_cache=True,
                store_kv=True,
            )

        num_transfer_tokens = self.get_num_transfer_tokens(block_length, denoising_steps)
        stop_tokens = self._prepare_stop_tokens(stopping_criteria, tokenizer, inputs_embeds.device)
        global_step = 0
        found_stop_token = False
        stop_pos = -1
        stop_chunk_end = -1
        start_time = time.perf_counter()

        for num_blocks in range(gen_blocks):
            block_start = prompt_length + num_blocks * block_length
            block_end = prompt_length + (num_blocks + 1) * block_length
            cur_x = x[:, block_start:block_end]
            cur_x_embeds = x_embeds[:, block_start:block_end, :]
            cur_step_map = step_map[:, block_start:block_end]
            cur_step_time = step_time[:, block_start:block_end]
            cur_attn_mask = full_attn_mask[block_start:block_end, :block_end].unsqueeze(0).unsqueeze(0)
            cur_position_ids = position_ids[:, block_start:block_end]

            for step in range(denoising_steps + 1):
                mask_index = (cur_x_embeds == mask_embeds).all(dim=-1)
                if mask_index.sum() == 0:
                    self(
                        inputs_embeds=cur_x_embeds,
                        attention_mask=cur_attn_mask,
                        position_ids=cur_position_ids,
                        past_key_values=past_key_values,
                        use_cache=True,
                        store_kv=True,
                    )
                    break
                outputs = self(
                    inputs_embeds=cur_x_embeds,
                    attention_mask=cur_attn_mask,
                    position_ids=cur_position_ids,
                    past_key_values=past_key_values,
                    use_cache=True,
                    store_kv=False,
                )
                x0, x0_p = self.sample_with_temperature_topk_topp(
                    outputs.logits,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                )
                confidence = torch.where(mask_index, x0_p, -torch.inf)
                transfer_index = self._select_transfer_index(
                    confidence,
                    dynamic_threshold,
                    int(num_transfer_tokens[min(step, denoising_steps - 1)].item()),
                )

                global_step += 1
                x0_embeds = self.get_input_embeddings()(x0)
                cur_x_embeds[transfer_index] = x0_embeds[transfer_index]
                cur_x[transfer_index] = x0[transfer_index]
                cur_step_map[transfer_index] = global_step
                cur_step_time[transfer_index] = time.perf_counter() - start_time

            if stop_tokens:
                generated = x[0, prompt_length:block_end]
                stop_offset = self._find_stop_position(generated, stop_tokens)
                if stop_offset is not None:
                    found_stop_token = True
                    stop_pos = prompt_length + stop_offset
                    stop_chunk_end = block_end
            if found_stop_token:
                break

        if found_stop_token:
            x = x[:, :stop_pos]
            step_map = step_map[:, :stop_chunk_end]
            step_time = step_time[:, :stop_chunk_end]

        return x[:, prompt_length:], step_map[:, prompt_length:], step_time[:, prompt_length:]


class MinerUDiffusionForConditionalGeneration(PreTrainedModel, GenerationMixin):
    config_class = MinerUDiffusionConfig

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, **kwargs):
        official_kwargs = dict(kwargs)
        official_kwargs.pop("trust_remote_code", None)
        if "dtype" in official_kwargs and "torch_dtype" not in official_kwargs:
            official_kwargs["torch_dtype"] = official_kwargs.pop("dtype")
        device = official_kwargs.pop("device", None)
        if device is not None and "device_map" not in official_kwargs:
            official_kwargs["device_map"] = device
            official_kwargs.setdefault("low_cpu_mem_usage", True)

        return super().from_pretrained(
            pretrained_model_name_or_path,
            *model_args,
            config=config,
            **official_kwargs,
        )

    def _init_weights(self, module):
        return

    def __init__(self, config: MinerUDiffusionConfig):
        super().__init__(config)
        if config.vision_model_type != "qwen2_vl":
            raise ValueError(f"Only qwen2_vl vision towers are supported, got: {config.vision_model_type}")
        self.vision_model = Qwen2VisionTransformerPretrainedModel._from_config(config.vision_model_config)
        self.vision_model.merger = nn.Identity()
        self.vision_abstractor = PerceiverProjection(
            projection_type=config.vision_projector_type,
            in_dim=config.vision_model_config.embed_dim,
            out_dim=config.language_model_config.hidden_size,
        )
        self.language_model = SDARForCausalLM(config.language_model_config)
        self.post_init()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def _prepare_inputs_embeds(
        self,
        input_ids: torch.LongTensor,
        pixel_values: Optional[torch.Tensor] = None,
        image_grid_thw: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        image_features = None
        if pixel_values is not None:
            image_features = self.get_image_features(pixel_values, image_grid_thw)
        return self._merge_input_and_image_features(input_ids, image_features)

    def _extract_vision_hidden_states(self, vision_outputs):
        if hasattr(vision_outputs, "last_hidden_state"):
            return vision_outputs.last_hidden_state
        if isinstance(vision_outputs, (tuple, list)):
            return vision_outputs[0]
        return vision_outputs

    def get_image_features(self, pixel_values, image_grid_thw):
        vision_outputs = self.vision_model(pixel_values, image_grid_thw)
        vision_hidden_states = self._extract_vision_hidden_states(vision_outputs)
        return self.vision_abstractor(vision_hidden_states)

    def _merge_input_and_image_features(self, input_ids, image_features):
        inputs_embeds = self.get_input_embeddings()(input_ids)
        if image_features is None:
            return inputs_embeds
        vision_mask = input_ids == self.config.image_token_id
        num_image_tokens = torch.count_nonzero(vision_mask).item()
        num_image_features = image_features.shape[:-1].numel()
        if num_image_tokens != num_image_features:
            raise ValueError(
                f"vision token count mismatch: {num_image_tokens} vs {num_image_features}"
            )
        return torch.masked_scatter(
            inputs_embeds,
            vision_mask.unsqueeze(-1),
            image_features.to(inputs_embeds.dtype).view(-1, image_features.size(-1)),
        )

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.BoolTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        image_grid_thw: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = None,
        return_dict: bool = True,
        **kwargs,
    ):
        inputs_embeds = self._prepare_inputs_embeds(input_ids, pixel_values, image_grid_thw)
        return self.language_model(
            input_ids=None,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            return_dict=return_dict,
            **kwargs,
        )

    @torch.no_grad()
    def generate(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.Tensor] = None,
        input_ids: Optional[torch.LongTensor] = None,
        **generate_kwargs,
    ):
        inputs_embeds = self._prepare_inputs_embeds(input_ids, pixel_values, image_grid_thw)
        return self.language_model.generate_with_embeds(inputs_embeds=inputs_embeds, **generate_kwargs)


MinerUDiffusion = MinerUDiffusionForConditionalGeneration


__all__ = ["MinerUDiffusionForConditionalGeneration", "MinerUDiffusion", "MinerUDiffusionConfig", "SDARConfig"]