reference/code/TAD/eval/model/generation_utils_dream.py

# Copyright 2025 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

# Modified from Dream repos: https://github.com/HKUNLP/Dream

import time
import warnings
import copy
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union

import torch
import torch.distributions as dists
from torch.nn import functional as F
from transformers import __version__
from transformers.generation.configuration_utils import (
    GenerationConfig
)
from transformers.utils import (
    ModelOutput,
    is_torchdynamo_compiling,
    logging,
)

logger = logging.get_logger(__name__)


def top_p_logits(logits, top_p=None):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    # Shift the indices to the right to keep the first token above the threshold
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
    return logits

def top_k_logits(logits, top_k=None):
    top_k = min(top_k, logits.size(-1))  # Safety check
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
    return logits


def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):

    if temperature > 0:
        logits = logits / temperature
    if top_p is not None and top_p < 1:
        logits = top_p_logits(logits, top_p)
    if top_k is not None:
        logits = top_k_logits(logits, top_k)
    probs = torch.softmax(logits, dim=-1)

    if temperature > 0:
        try:
            x0 = dists.Categorical(probs=probs).sample()
            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
        except:
            confidence, x0 = probs.max(dim=-1)
    else:
        confidence, x0 = probs.max(dim=-1)
    
    if margin_confidence:
        sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
        # Extract top1 and top2 probabilities
        top1_probs = sorted_probs[:, 0] 
        top2_probs = sorted_probs[:, 1] 
        # Calculate confidence as top1 - top2
        confidence = top1_probs - top2_probs 
    
    if neg_entropy:
        epsilon = 1e-10
        log_probs = torch.log(probs + epsilon)
        confidence = torch.sum(probs * log_probs, dim=-1)
    
    return confidence, x0


def sample_tokens_with_entropy(logits, temperature=1.0):
    """Sample tokens and return their entropy values.
    
    Used by multi-block generation for entropy-based token selection.
    Returns (entropy, samples) where lower entropy = higher confidence.
    """
    original_probs = torch.softmax(logits, dim=-1)
    log_probs = torch.log(original_probs + 1e-8)
    entropy = -torch.sum(original_probs * log_probs, dim=-1)
    
    if temperature == 0:
        samples = torch.argmax(logits, dim=-1)
    else:
        scaled_logits = logits / temperature
        probs = torch.softmax(scaled_logits, dim=-1)
        samples = torch.multinomial(probs, num_samples=1).squeeze(-1)
    
    return entropy, samples


@dataclass
class DreamModelOutput(ModelOutput):
    sequences: torch.LongTensor = None
    history: Optional[Tuple[torch.FloatTensor]] = None
    # total number of forward iterations actually used during generation
    nfe: Optional[int] = None


class DreamGenerationConfig(GenerationConfig):
    def __init__(self, **kwargs):
        self.temperature: float = kwargs.pop("temperature", 0.0)
        self.top_p: Optional[float] = kwargs.pop("top_p", None)
        self.top_k: Optional[int] = kwargs.pop("top_k", None)
        self.max_length = kwargs.pop("max_length", 20)
        self.max_new_tokens = kwargs.pop("max_new_tokens", None)
        # diffusion specific params
        self.eps: float = kwargs.pop("eps", 1e-3)
        self.steps: int = kwargs.pop("steps", 512)
        self.alg: str = kwargs.pop("alg", 'origin')
        self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)

        # Parameters that define the output variables of `generate`
        self.num_return_sequences: int = kwargs.pop("num_return_sequences", 1)
        self.return_dict_in_generate: bool = kwargs.pop("return_dict_in_generate", False)
        self.output_history: bool = kwargs.pop("output_history", False)

        # Special tokens that can be used at generation time
        self.mask_token_id = kwargs.pop("mask_token_id", None)
        self.pad_token_id = kwargs.pop("pad_token_id", None)
        self.bos_token_id = kwargs.pop("bos_token_id", None)
        self.eos_token_id = kwargs.pop("eos_token_id", None)

        # Wild card
        self.generation_kwargs = kwargs.pop("generation_kwargs", {})

        # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub
        # interface.
        self._from_model_config = kwargs.pop("_from_model_config", False)
        self._commit_hash = kwargs.pop("_commit_hash", None)
        self.transformers_version = kwargs.pop("transformers_version", __version__)

        # Additional attributes without default values
        if not self._from_model_config:
            # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a
            # model's default configuration file
            for key, value in kwargs.items():
                try:
                    setattr(self, key, value)
                except AttributeError as err:
                    logger.error(f"Can't set {key} with value {value} for {self}")
                    raise err

        # Validate the values of the attributes
        self.validate(is_init=True)

    def validate(self, is_init=False):
        pass

class DreamGenerationMixin:
    @staticmethod
    def _expand_inputs_for_generation(
        expand_size: int = 1,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None
    ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
        """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
        # Do not call torch.repeat_interleave if expand_size is 1 because it clones
        # the input tensor and thus requires more memory although no change is applied
        if expand_size == 1:
            return input_ids, attention_mask
        if input_ids is not None:
            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
        if attention_mask is not None:
            attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)
        return input_ids, attention_mask

    def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
        """Performs validation related to the resulting generated length"""

        # Can't throw warnings/exceptions during compilation
        if is_torchdynamo_compiling():
            return

        # 1. Max length warnings related to poor parameterization
        if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
            # 20 is the default max_length of the generation config
            warnings.warn(
                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
                "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
                "generation.",
                UserWarning,
            )
        if input_ids_length >= generation_config.max_length:
            input_ids_string = "input_ids"
            raise ValueError(
                f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
                " increasing `max_length` or, better yet, setting `max_new_tokens`."
            )

    def _prepare_generated_length(
        self,
        generation_config,
        has_default_max_length,
        input_ids_length,
    ):
        """Prepared max and min length in generation configs to avoid clashes between similar attributes"""

        if generation_config.max_new_tokens is not None:
            if not has_default_max_length and generation_config.max_length is not None:
                logger.warning(
                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                    "Please refer to the documentation for more information. "
                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
                )
            generation_config.max_length = generation_config.max_new_tokens + input_ids_length

        elif has_default_max_length:
            if generation_config.max_length == DreamGenerationConfig().max_length:
                generation_config.max_length = generation_config.max_length + input_ids_length
                max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
                if max_position_embeddings is not None:
                    generation_config.max_length = min(generation_config.max_length, max_position_embeddings)

        return generation_config

    def _prepare_generation_config(
        self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict
    ) -> DreamGenerationConfig:
        """
        Prepares the base generation config, then applies any generation configuration options from kwargs. This
        function handles retrocompatibility with respect to configuration files.
        """
        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
        using_model_generation_config = False
        if generation_config is None:
            generation_config = DreamGenerationConfig.from_model_config(self.config)
            using_model_generation_config = True

        # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config`
        # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an
        # exception will be raised in `_validate_model_kwargs`
        if not is_torchdynamo_compiling():
            generation_config = copy.deepcopy(generation_config)
            _kwargs = generation_config.update(**kwargs)
            # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model
            if not using_model_generation_config:
                if generation_config.bos_token_id is None:
                    generation_config.bos_token_id = self.generation_config.bos_token_id
                if generation_config.eos_token_id is None:
                    generation_config.eos_token_id = self.generation_config.eos_token_id
                if generation_config.pad_token_id is None:
                    generation_config.pad_token_id = self.generation_config.pad_token_id
                if generation_config.mask_token_id is None:
                    generation_config.mask_token_id = self.generation_config.mask_token_id

        return generation_config

    def _prepare_special_tokens(
        self,
        generation_config: DreamGenerationConfig,
        device: Optional[Union[torch.device, str]] = None,
    ):
        """
        Prepares the special tokens for generation, overwriting the generation config with their processed versions
        converted to tensor.
        Note that `generation_config` is changed in place and stops being serializable after this method is called.
        That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the
        function). However, if called outside `generate`, consider creating a copy of `generation_config` first.
        """

        # Convert special tokens to tensors
        def _tensor_or_none(token, device=None):
            if token is None:
                return token

            device = device if device is not None else self.device
            if isinstance(token, torch.Tensor):
                return token.to(device)
            return torch.tensor(token, device=device, dtype=torch.long)

        bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
        eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
        pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
        mask_token_tensor = _tensor_or_none(generation_config.mask_token_id, device=device)

        # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
        if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
            eos_token_tensor = eos_token_tensor.unsqueeze(0)

        # Set pad token if unset (and there are conditions to do so)
        if pad_token_tensor is None and eos_token_tensor is not None:
            pad_token_tensor = eos_token_tensor[0]
            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")

        # Update generation config with the updated special tokens tensors
        # NOTE: this must be written into a different attribute name than the one holding the original special tokens
        # (in their non-tensor form), in order to enable end-to-end compilation. See
        # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations
        generation_config._bos_token_tensor = bos_token_tensor
        generation_config._eos_token_tensor = eos_token_tensor
        generation_config._pad_token_tensor = pad_token_tensor
        generation_config._mask_token_tensor = mask_token_tensor

    @torch.no_grad()
    def diffusion_generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        generation_config: Optional[DreamGenerationConfig] = None,
        **kwargs,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
        generation_config = self._prepare_generation_config(generation_config, **kwargs)
        generation_tokens_hook_func = kwargs.pop("generation_tokens_hook_func", lambda step, x, logits: x)
        generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)

        # 2. Define model inputs
        assert inputs is not None
        input_ids = inputs
        device = input_ids.device
        attention_mask = kwargs.pop("attention_mask", None)
        self._prepare_special_tokens(generation_config, device=device)

        # 3. Prepare `max_length`.
        input_ids_length = input_ids.shape[-1]
        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        generation_config = self._prepare_generated_length(
            generation_config=generation_config,
            has_default_max_length=has_default_max_length,
            input_ids_length=input_ids_length,
        )

        self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
        
        # 4. Check input_ids
        if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
            warnings.warn(
                "You are calling .generate() with the `input_ids` being on a device type different"
                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
                " Please make sure that you have put `input_ids` to the"
                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
                " running `.generate()`.",
                UserWarning,
            )
        if (
            hasattr(generation_config, "pad_token_id") and
            torch.any(input_ids == generation_config.pad_token_id) and 
            attention_mask is None
        ):
            warnings.warn(
                "Padding was detected but no attention mask is passed here. For correct "
                "generation results, please set `attention_mask` when batch-padding inputs.",
                UserWarning,
            )

        input_ids, attention_mask = self._expand_inputs_for_generation(
            expand_size=generation_config.num_return_sequences,
            input_ids=input_ids,
            attention_mask=attention_mask 
        )
        threshold = kwargs.get("threshold", 0.9)

        # Support block-wise generation even without cache
        block_length = kwargs.get("block_length", None)
        early_stop = kwargs.get("early_stop", False)

        result = self._sample(
            input_ids,
            attention_mask=attention_mask,
            generation_config=generation_config,
            generation_tokens_hook_func=generation_tokens_hook_func,
            generation_logits_hook_func=generation_logits_hook_func,
            threshold=threshold,
            block_length=block_length,
            early_stop=early_stop,
        )
        return result

    def _sample(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor],
        generation_config: DreamGenerationConfig,
        generation_tokens_hook_func,
        generation_logits_hook_func,
        threshold: Optional[float] = 0.9,
        block_length: Optional[int] = None,
        early_stop: bool = False,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        # init values
        output_history = generation_config.output_history
        return_dict_in_generate = generation_config.return_dict_in_generate
        max_length = generation_config.max_length
        mask_token_id = generation_config.mask_token_id
        steps = generation_config.steps
        eps = generation_config.eps
        alg = generation_config.alg
        alg_temp = generation_config.alg_temp
        temperature = generation_config.temperature
        top_p = generation_config.top_p
        top_k = generation_config.top_k
        eos_token_id = generation_config.eos_token_id if early_stop else None

        histories = [] if (return_dict_in_generate and output_history) else None
        start_time = time.time()
        # pad input_ids to max_length
        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id)
        gen_length = max_length - input_ids.shape[1]

        # prepare attention mask/tok_idx
        if attention_mask is not None and torch.any(attention_mask == 0.0):
            attention_mask = F.pad(attention_mask, (0, max_length - attention_mask.shape[1]), value=1.0)
            tok_idx = attention_mask.long().cumsum(-1) - 1
            tok_idx.masked_fill_(attention_mask == 0, 1)
            attention_mask = torch.logical_and(
                attention_mask.unsqueeze(1).unsqueeze(-2),
                attention_mask.unsqueeze(1).unsqueeze(-1),
            )
        else:
            tok_idx = None
            attention_mask = "full"

        # Determine block config: default to single block (legacy behavior)
        if block_length is None:
            block_length = gen_length if gen_length > 0 else 1
        assert gen_length % block_length == 0, f"gen_length ({gen_length}) must be divisible by block_length ({block_length})"
        num_blocks = max(gen_length // block_length, 1)

        assert steps % num_blocks == 0, f"steps ({steps}) must be divisible by num_blocks ({num_blocks})"
        steps_per_block = steps // num_blocks

        # per-block timesteps
        timesteps_block = torch.linspace(1, eps, steps_per_block + 1, device=x.device)

        # initial hook
        x = generation_tokens_hook_func(None, x, None)

        # iterate blocks without cache
        total_nfe = 0
        prompt_length = input_ids.shape[1]
        has_eos = False
        for num_block in range(num_blocks):
            # Early stop: skip remaining blocks if EOS already found
            if eos_token_id is not None and has_eos:
                break
            current_block_start = input_ids.shape[1] + num_block * block_length
            current_block_end = current_block_start + block_length

            i = 0
            while True:
                # stop if current block completed
                if (x[:, current_block_start:current_block_end] == mask_token_id).sum() == 0:
                    break

                # forward
                logits = self(x, attention_mask, tok_idx).logits
                logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)
                total_nfe += 1

                # logits hook
                logits = generation_logits_hook_func(i, x, logits)

                # build mask over full sequence but restrict to current block
                mask_index_full = (x == mask_token_id)
                # zero out positions outside current block on and after start
                mask_index_full[:, :current_block_start] = False
                mask_index_full[:, current_block_end:] = False

                mask_logits = logits[mask_index_full]

                if alg == 'confidence_threshold':
                    confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
                    x_block = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
                    x_block[mask_index_full] = x0.clone()
                    full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
                    full_confidence[mask_index_full] = confidence
                    # strictly keep within current block
                    full_confidence[:, :current_block_start] = -torch.inf
                    full_confidence[:, current_block_end:] = -torch.inf

                    current_transfer_tokens = (x[:, current_block_start:current_block_end] == mask_token_id).sum()
                    selected_confidence, select_index = torch.topk(full_confidence, current_transfer_tokens)
                    transfer_index = torch.zeros_like(x, device=x.device, dtype=torch.bool)
                    select_index = select_index.to(x.device)
                    transfer_index[0, select_index[0]] = True
                    for k in range(1, current_transfer_tokens):
                        if selected_confidence[0, k] < threshold:
                            transfer_index[0, select_index[0, k]] = False
                    x[transfer_index] = x_block[transfer_index].clone()
                elif alg == 'entropy_threshold':
                    # entropy_threshold: decode tokens with entropy below threshold
                    # Lower entropy = more confident, so we decode low-entropy tokens first
                    neg_entropy, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
                    # neg_entropy is negative entropy, so higher neg_entropy = lower actual entropy = more confident
                    entropy = -neg_entropy  # convert to actual entropy (positive value, lower = more confident)
                    
                    x_block = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
                    x_block[mask_index_full] = x0.clone()
                    full_entropy = torch.full_like(x, torch.inf, device=self.device, dtype=logits.dtype)
                    full_entropy[mask_index_full] = entropy
                    # strictly keep within current block
                    full_entropy[:, :current_block_start] = torch.inf
                    full_entropy[:, current_block_end:] = torch.inf

                    current_transfer_tokens = (x[:, current_block_start:current_block_end] == mask_token_id).sum()
                    # Sort by entropy ascending (lowest entropy first = most confident first)
                    selected_entropy, select_index = torch.topk(full_entropy, current_transfer_tokens, largest=False)
                    transfer_index = torch.zeros_like(x, device=x.device, dtype=torch.bool)
                    select_index = select_index.to(x.device)
                    transfer_index[0, select_index[0]] = True
                    for k in range(1, current_transfer_tokens):
                        if selected_entropy[0, k] > threshold:
                            # entropy > threshold means uncertain, don't decode
                            transfer_index[0, select_index[0, k]] = False
                    x[transfer_index] = x_block[transfer_index].clone()
                else:
                    # diffusion style transfer count within block
                    if i == steps_per_block:
                        break
                    t = timesteps_block[i]
                    s = timesteps_block[i + 1]

                    # compute confidence for the masked positions within block
                    if alg == 'maskgit_plus':
                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
                    elif alg == 'topk_margin':
                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True)
                    elif alg == 'entropy':
                        confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
                    elif alg == 'origin':
                        # origin behaves like proportional transfer probability
                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
                    else:
                        raise RuntimeError(f"Unknown alg: {alg}")

                    # full confidence tensor limited to current block
                    full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
                    full_confidence[mask_index_full] = confidence
                    full_confidence[:, :current_block_start] = -torch.inf
                    full_confidence[:, current_block_end:] = -torch.inf

                    # compute transfer tokens count in block
                    num_mask_token = mask_index_full.sum() / mask_index_full.shape[0]
                    number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps_per_block - 1 else int(num_mask_token)

                    if number_transfer_tokens > 0:
                        if alg_temp is None or alg_temp == 0:
                            _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)
                        else:
                            full_confidence = full_confidence / alg_temp
                            full_confidence = F.softmax(full_confidence, dim=-1)
                            transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
                        x_block = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
                        x_block[mask_index_full] = x0.clone()
                        row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)
                        x[row_indices, transfer_index] = x_block[row_indices, transfer_index]
                    i += 1

                # tokens hook and history
                x = generation_tokens_hook_func(i, x, logits)
                if histories is not None:
                    histories.append(x.clone())

                # Early stop: check for EOS after each decode step
                if eos_token_id is not None:
                    gen_region = x[:, prompt_length:]
                    eos_found = (gen_region == eos_token_id) & (gen_region != mask_token_id)
                    if eos_found.any():
                        pos = torch.arange(gen_region.shape[1], device=x.device).unsqueeze(0)
                        first_eos_rel = torch.where(eos_found, pos, gen_region.shape[1]).amin(dim=1)
                        first_eos_abs = prompt_length + first_eos_rel[0].item()
                        x[:, first_eos_abs + 1:] = eos_token_id
                        has_eos = True
                        break
        
        end_time = time.time()
        if return_dict_in_generate:
            return DreamModelOutput(
                sequences=x,
                history=histories,
                nfe=total_nfe,
            )
        else:
            return x

    @torch.no_grad()
    def generate_multi_block(
        self,
        inputs: Optional[torch.Tensor] = None,
        generation_config: Optional[DreamGenerationConfig] = None,
        **kwargs,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        """
        Entry point for multi-block pipelined parallel decoding.
        Mirrors diffusion_generate() but dispatches to _sample_multi_block().
        """
        # 1. Handle generation_config and kwargs
        generation_config = self._prepare_generation_config(generation_config, **kwargs)

        # 2. Define model inputs
        assert inputs is not None
        input_ids = inputs
        device = input_ids.device
        attention_mask = kwargs.pop("attention_mask", None)
        self._prepare_special_tokens(generation_config, device=device)

        # 3. Prepare max_length
        input_ids_length = input_ids.shape[-1]
        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        generation_config = self._prepare_generated_length(
            generation_config=generation_config,
            has_default_max_length=has_default_max_length,
            input_ids_length=input_ids_length,
        )
        self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)

        # 4. Check device
        if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
            warnings.warn(
                "You are calling .generate() with the `input_ids` being on a device type different"
                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation.",
                UserWarning,
            )
        if (
            hasattr(generation_config, "pad_token_id") and
            torch.any(input_ids == generation_config.pad_token_id) and
            attention_mask is None
        ):
            warnings.warn(
                "Padding was detected but no attention mask is passed here. For correct "
                "generation results, please set `attention_mask` when batch-padding inputs.",
                UserWarning,
            )

        input_ids, attention_mask = self._expand_inputs_for_generation(
            expand_size=generation_config.num_return_sequences,
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

        # Extract multi-block specific parameters
        threshold = kwargs.get("threshold", 0.9)
        block_size = kwargs.get("block_size", 32)
        block_add_threshold = kwargs.get("block_add_threshold", 0.5)
        decoded_token_threshold = kwargs.get("decoded_token_threshold", 0.5)
        early_stop = kwargs.get("early_stop", False)

        result = self._sample_multi_block(
            input_ids,
            attention_mask=attention_mask,
            generation_config=generation_config,
            threshold=threshold,
            block_size=block_size,
            block_add_threshold=block_add_threshold,
            decoded_token_threshold=decoded_token_threshold,
            early_stop=early_stop,
        )
        return result

    def _sample_multi_block(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor],
        generation_config: DreamGenerationConfig,
        threshold: float = 0.9,
        block_size: int = 32,
        block_add_threshold: float = 0.5,
        decoded_token_threshold: float = 0.5,
        early_stop: bool = False,
    ) -> Union[DreamModelOutput, torch.LongTensor]:
        """
        Pipelined parallel decoding without cache.

        Args:
            block_add_threshold: Add new block when last block progress >= this threshold.
                                Set to 1.0 for fully sequential processing.
            decoded_token_threshold: Block becomes fully activated when previous block progress >= this threshold.
                                    Set to 1.0 for fully sequential processing.
            threshold: Entropy threshold for decoding (lower entropy = higher confidence).

        When block_add_threshold=1.0 and decoded_token_threshold=1.0, this method behaves
        identically to standard generation with sequential block processing.
        """
        return_dict_in_generate = generation_config.return_dict_in_generate
        max_length = generation_config.max_length
        mask_token_id = generation_config.mask_token_id
        temperature = generation_config.temperature
        alg = generation_config.alg
        eos_token_id = generation_config.eos_token_id if early_stop else None

        max_new_tokens = max_length - input_ids.shape[1]
        prompt_length = input_ids.shape[1]
        x = F.pad(input_ids, (0, max_new_tokens), value=mask_token_id)

        # Prepare attention mask
        if attention_mask is not None and torch.any(attention_mask == 0.0):
            attention_mask_padded = F.pad(attention_mask, (0, max_new_tokens), value=1.0)
            tok_idx = attention_mask_padded.long().cumsum(-1) - 1
            tok_idx.masked_fill_(attention_mask_padded == 0, 1)
            attn_mask_4d = torch.logical_and(
                attention_mask_padded.unsqueeze(1).unsqueeze(-2),
                attention_mask_padded.unsqueeze(1).unsqueeze(-1),
            )
        else:
            tok_idx = None
            attn_mask_4d = "full"

        # Track block states: {block_id: {start, end, mask_count, total_masks, is_complete}}
        # Initialize with prompt block
        block_states = {
            0: {
                "start": 0,
                "end": input_ids.shape[1],
                "mask_count": 0,
                "total_masks": input_ids.shape[1],
                "is_complete": True,
            }
        }

        # Create first generation block
        num_blocks = max_new_tokens // block_size
        next_block_id = 1
        if next_block_id <= num_blocks:
            block_start = input_ids.shape[1] + (next_block_id - 1) * block_size
            block_end = min(block_start + block_size, input_ids.shape[1] + max_new_tokens)
            should_activate = 1.0 >= decoded_token_threshold  # prompt progress is always 1.0
            block_states[next_block_id] = {
                "start": block_start,
                "end": block_end,
                "mask_count": block_end - block_start,
                "total_masks": block_end - block_start,
                "is_complete": should_activate,
            }
            next_block_id += 1

        nfe = 0

        while True:
            # Check if all blocks are exhausted AND no more blocks to create
            mask_index = x == mask_token_id
            total_masks = mask_index[:, prompt_length:].sum()

            if total_masks == 0 and next_block_id > num_blocks:
                break

            nfe += 1

            # Early stop: check for EOS token
            if eos_token_id is not None:
                gen_region = x[:, prompt_length:]
                eos_found = (gen_region == eos_token_id) & (gen_region != mask_token_id)
                if eos_found.any():
                    pos = torch.arange(gen_region.shape[1], device=x.device).unsqueeze(0)
                    first_eos_rel = torch.where(eos_found, pos, gen_region.shape[1]).amin(dim=1)
                    first_eos_abs = prompt_length + first_eos_rel[0].item()
                    x[:, first_eos_abs + 1:] = eos_token_id
                    # Create all remaining blocks after EOS and mark them as complete
                    while next_block_id <= num_blocks:
                        bs = prompt_length + (next_block_id - 1) * block_size
                        be = min(bs + block_size, prompt_length + max_new_tokens)
                        if bs > first_eos_abs:
                            block_states[next_block_id] = {
                                "start": bs, "end": be, "mask_count": 0,
                                "total_masks": be - bs, "is_complete": True,
                            }
                            next_block_id += 1
                        else:
                            break
                    if (x == mask_token_id)[:, prompt_length:].sum() == 0:
                        break

            # Update block activation states
            def update_block_activation_states():
                """Update which blocks should be fully activated based on previous block progress."""
                for bid in sorted(block_states.keys()):
                    if bid > 0 and not block_states[bid]["is_complete"]:
                        prev_progress = (
                            1 - block_states[bid - 1]["mask_count"] / block_states[bid - 1]["total_masks"]
                        )
                        if prev_progress >= decoded_token_threshold:
                            block_states[bid]["is_complete"] = True

            update_block_activation_states()

            # Add new block dynamically based on last block's progress
            if next_block_id <= num_blocks:
                last_bid = max(block_states.keys())
                if last_bid > 0:  # Not just prompt
                    last_progress = (
                        1 - block_states[last_bid]["mask_count"] / block_states[last_bid]["total_masks"]
                    )
                    should_add_block = (last_progress >= block_add_threshold) or (block_states[last_bid]["mask_count"] == 0)

                    if should_add_block:
                        block_start = input_ids.shape[1] + (next_block_id - 1) * block_size
                        block_end = min(block_start + block_size, input_ids.shape[1] + max_new_tokens)
                        if block_end > block_start:
                            actual_mask_count = (x[:, block_start:block_end] == mask_token_id).sum().item()
                            prev_bid = next_block_id - 1
                            prev_progress = (
                                1 - block_states[prev_bid]["mask_count"] / block_states[prev_bid]["total_masks"]
                            )
                            should_activate = prev_progress >= decoded_token_threshold

                            block_states[next_block_id] = {
                                "start": block_start,
                                "end": block_end,
                                "mask_count": actual_mask_count,
                                "total_masks": block_end - block_start,
                                "is_complete": should_activate,
                            }
                            next_block_id += 1

            # Find the rightmost block that should be processed
            rightmost_active_bid = 0
            for bid in sorted(block_states.keys()):
                if block_states[bid]["is_complete"] or block_states[bid]["mask_count"] > 0:
                    rightmost_active_bid = bid

            if rightmost_active_bid == 0:
                break

            active_end = block_states[rightmost_active_bid]["end"]

            # Forward pass on entire sequence
            model_output = self(x, attn_mask_4d, tok_idx)
            logits = model_output.logits
            logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)

            # Mask out future blocks (positions after active_end)
            mask_index_for_decode = mask_index.clone()
            mask_index_for_decode[:, active_end:] = 0

            # Decode with entropy-based threshold
            if alg == 'entropy_threshold':
                mask_logits = logits[mask_index_for_decode]

                entropy, x0 = sample_tokens_with_entropy(mask_logits, temperature=temperature)

                x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
                full_entropy = torch.full_like(x, torch.inf, device=self.device, dtype=logits.dtype)

                x_[mask_index_for_decode] = x0.clone()
                full_entropy[mask_index_for_decode] = entropy

                current_transfer_tokens = mask_index_for_decode.sum()

                selected_entropy, select_index = torch.topk(full_entropy, current_transfer_tokens, largest=False)
                transfer_index = torch.zeros_like(x, device=x.device, dtype=torch.bool)

                select_index = select_index.to(x.device)
                transfer_index[0, select_index[0]] = True
                for k in range(1, current_transfer_tokens):
                    if selected_entropy[0, k] < threshold:
                        transfer_index[0, select_index[0, k]] = True
                    else:
                        transfer_index[0, select_index[0, k]] = False

                # For fully activated blocks, ensure at least one token is decoded (guaranteed progress)
                first_fully_activated_bid = None
                for bid in sorted(block_states.keys()):
                    if bid > 0 and block_states[bid]["is_complete"] and block_states[bid]["mask_count"] > 0:
                        first_fully_activated_bid = bid
                        break

                if first_fully_activated_bid is not None:
                    start, end = block_states[first_fully_activated_bid]["start"], block_states[first_fully_activated_bid]["end"]
                    block_transfer = transfer_index[:, start:end]

                    if not block_transfer.any():
                        # Force decode the lowest entropy token in this fully activated block
                        block_mask = mask_index_for_decode[:, start:end]
                        block_entropy = full_entropy[:, start:end]
                        block_entropy = torch.where(block_mask, block_entropy, torch.inf)
                        best_idx = block_entropy[0].argmin()
                        transfer_index[0, start + best_idx] = True

                # Apply the decoded tokens
                x[transfer_index] = x_[transfer_index].clone()

                # Update block states based on which positions were decoded
                for bid in sorted(block_states.keys()):
                    if bid > 0 and block_states[bid]["mask_count"] > 0:
                        start, end = block_states[bid]["start"], block_states[bid]["end"]
                        block_decoded = transfer_index[:, start:end].sum().item()
                        if block_decoded > 0:
                            block_states[bid]["mask_count"] -= block_decoded

            if nfe > 10000:
                break

        if return_dict_in_generate:
            return DreamModelOutput(sequences=x, nfe=nfe), nfe
        return x, nfe