Update generation_utils.py

generation_utils.py

@@ -30,23 +30,8 @@ from transformers.utils import (
     is_torchdynamo_compiling,
     logging,
 )
-from .generate_from_llada import  get_num_transfer_tokens_sch
 logger = logging.get_logger(__name__)
-
-import sys
-import pdb
-class ForkedPdb(pdb.Pdb):
-    """
-    PDB Subclass for debugging multi-processed code
-    Suggested in: https://stackoverflow.com/questions/4716533/how-to-attach-debugger-to-a-python-subproccess
-    """
-    def interaction(self, *args, **kwargs):
-        _stdin = sys.stdin
-        try:
-            sys.stdin = open('/dev/stdin')
-            pdb.Pdb.interaction(self, *args, **kwargs)
-        finally:
-            sys.stdin = _stdin
+from tqdm import tqdm
 
 
 def top_p_logits(logits, top_p=None):
@@ -70,123 +55,59 @@ def top_k_logits(logits, top_k=None):
     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(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):
-    """
-    从给定的 logits 中采样或贪心选取 token，并返回置信度和 token ID。
-
-    参数：
-        logits (Tensor)：形状 [batch_size, vocab_size]，模型对各候选 token 的打分（未经 softmax）。
-        temperature (float)：温度系数，默认 0.0。>0 时按概率采样，=0 时贪心选取。
-        top_p (float 或 None)：核采样参数（nucleus sampling），若指定且 <1，只保留累计概率前 top_p 的 token。
-        top_k (int 或 None)：前 k 采样参数（top-k sampling），若指定，只从概率最高的 k 个 token 中选取。
-        margin_confidence (bool)：是否使用 top1−top2 之差作为置信度，默认 False。
-        neg_entropy (bool)：是否使用负熵（−∑p·logp）作为置信度，默认 False。
-
-    返回：
-        confidence (Tensor)：形状 [batch_size] 的置信度值（可用概率、margin 差值或负熵）。
-        x0 (Tensor)：形状 [batch_size] 的 int64 张量，表示采样或贪心得到的 token ID。
-    """
-
-    # ======================================================
-    # 1. 温度缩放 (Temperature Scaling)
-    # ======================================================
+def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False, 
+                  repeat_penalty=1.0, max_position_penalty=1.0, past_x=None, mask_id=None,):
     if temperature > 0:
-        # 当 temperature>0 时，将 logits 除以 temperature，使得 softmax 分布更平滑或更尖锐
         logits = logits / temperature
 
-    # ======================================================
-    # 2. Top-p (Nucleus) 与 Top-k 过滤
-    # ======================================================
     if top_p is not None and top_p < 1:
-        # 调用 top_p_logits，保留累计概率达到 top_p 的 token，其它 logits 置为很小的负值
         logits = top_p_logits(logits, top_p)
     if top_k is not None:
-        # 调用 top_k_logits，仅保留概率最高的 top_k 个 token，其它 logits 置为很小的负值
         logits = top_k_logits(logits, top_k)
 
-    # ======================================================
-    # 3. 计算概率分布 (Softmax)
-    # ======================================================
-    probs = torch.softmax(logits, dim=-1)
-    # 此时 probs 形状为 [batch_size, vocab_size]，每行和为 1
+    if repeat_penalty != 1.0:
+        select_mask = torch.logical_and((past_x != 0), (past_x != mask_id))
+        generated_tokens = set(past_x[select_mask].tolist())
+        for token in set(generated_tokens):
+            logits[:, token][logits[:, token] < 0] *= repeat_penalty
+            logits[:, token][logits[:, token] >= 0] /= repeat_penalty
+
+    if max_position_penalty != 1.0:
+        token_length = logits.shape[-2]
+        if token_length > 100:
+            penalty_map = [i / (token_length - 100) * (max_position_penalty - 1.0) + 1.0 
+                           for i in range(token_length - 100)]
+            penalty_map = torch.tensor(penalty_map).unsqueeze(-1).to(logits.device).to(logits.dtype)
+            penalty_map = torch.cat([torch.ones_like(logits[:100, :1]), penalty_map], dim=0)
+            penalty_map = penalty_map.repeat(1, logits.shape[-1])
+        
+            logits[logits < 0] *= penalty_map[logits < 0]
+            logits[logits >= 0] /= penalty_map[logits >= 0]
 
-    # ======================================================
-    # 4. 根据 temperature 决定采样或贪心选取
-    # ======================================================
+    probs = torch.softmax(logits, dim=-1)
     if temperature > 0:
-        # 随机采样分支：从 Categorical 分布中采样 token
         try:
-            # 从多项分布中采样得到 token ID，形状 [batch_size]
             x0 = dists.Categorical(probs=probs).sample()
-            # 用 gather 取出对应位置的概率值作为置信度，形状 [batch_size]
             confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
         except:
-            # 若采样出错（如概率分布不合法），退化为贪心选取
             confidence, x0 = probs.max(dim=-1)
     else:
-        # 当 temperature=0 时，直接贪心选取概率最大的 token
         confidence, x0 = probs.max(dim=-1)
 
-    # ======================================================
-    # 5. margin_confidence: 使用 top1−top2 差值作为置信度
-    # ======================================================
     if margin_confidence:
-        # 将每行概率按降序排序，sorted_probs[:,0] 为 top1，sorted_probs[:,1] 为 top2
         sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
         top1_probs = sorted_probs[:, 0]
         top2_probs = sorted_probs[:, 1]
-        # 置信度设为 top1_probs − top2_probs
         confidence = top1_probs - top2_probs
 
-    # ======================================================
-    # 6. neg_entropy: 使用负熵（−∑ p·log p）作为置信度
-    # ======================================================
     if neg_entropy:
         epsilon = 1e-10
-        # 为避免 log(0) 产生 −inf，加上一个小常数 epsilon
         log_probs = torch.log(probs + epsilon)
-        # 计算 ∑ p_i * log p_i，结果是负熵值（值越接近 0，表示分布更“尖锐”）
         confidence = torch.sum(probs * log_probs, dim=-1)
 
     return confidence, x0
 
 
-
 @dataclass
 class DreamModelOutput(ModelOutput):
     sequences: torch.LongTensor = None
@@ -398,6 +319,10 @@ class DreamGenerationMixin:
         generation_config: Optional[DreamGenerationConfig] = None,
         inputs_embeds=None,
         prefix_lm=False,
+        alg=None,
+        block_size=-1,
+        cfg=0.0,
+        add_boa_token=False,
         **kwargs,
     ) -> Union[DreamModelOutput, torch.LongTensor]:
         # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
@@ -406,7 +331,6 @@ class DreamGenerationMixin:
         generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)
         # breakpoint()
         # 2. Define model inputs
-        # import pdb;pdb.set_trace()
         if inputs is not None:
             input_ids = inputs
             device = input_ids.device
@@ -440,7 +364,6 @@ class DreamGenerationMixin:
                 f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
                 " running `.generate()`.",
                 UserWarning,
-            
             )
         # breakpoint()
         if (
@@ -454,14 +377,13 @@ class DreamGenerationMixin:
                 "generation results, please set `attention_mask` when batch-padding inputs.",
                 UserWarning,
             )
-        assert generation_config.num_return_sequences == 1, "Currently, we only support num_return_sequences = 1 for diffusion generation."
-        # import pdb;pdb.set_trace()
+        assert generation_config.num_return_sequences == 1, \
+            "Currently, we only support num_return_sequences = 1 for diffusion generation."
         input_ids, attention_mask = self._expand_inputs_for_generation(
             expand_size=generation_config.num_return_sequences,
             input_ids=input_ids,
             attention_mask=attention_mask 
         )
-
         result = self._sample(
             input_ids,
             attention_mask=attention_mask,
@@ -471,9 +393,14 @@ class DreamGenerationMixin:
             inputs_embeds=inputs_embeds,
             device=device,
             prefix_lm=prefix_lm,
+            alg=alg,
+            block_size=block_size,
+            cfg=cfg,
+            add_boa_token=add_boa_token,
             **kwargs,
         )
         return result
+
     def _sample(
         self,
         input_ids: torch.LongTensor,
@@ -484,223 +411,213 @@ class DreamGenerationMixin:
         inputs_embeds=None,
         prefix_lm=False,
         device=None,
-        schedule_kwargs=None,
-        schedule=None,
         step_ratio=None,
+        penalty=1.2,
+        alg=None,
+        block_size=None,
+        add_boa_token=False,
+        max_position_penalty=1.0,
+        repeat_penalty=1.0,
+        cfg=0.0,
         **kwargs,
     ) -> Union[DreamModelOutput, torch.LongTensor]:
-        # 1. 从 generation_config 中提取常用参数
-        output_history = generation_config.output_history            # 是否保存每一步的中间结果
-        # output_history = True
-        return_dict_in_generate = generation_config.return_dict_in_generate  # 生成时是否返回字典形式
-        max_length = generation_config.max_length                     # 生成后序列的最大长度（包括前缀）
-        mask_token_id = generation_config.mask_token_id               # [MASK] 的 token ID
-        max_new_tokens = generation_config.max_new_tokens             # 最多新增的 token 数量
-        steps = min(generation_config.steps, max_new_tokens)          # 实际去噪步数，不能超过最大新增 token 数
-        eps = generation_config.eps                                    # 噪声下限，用于时刻表
-        alg = generation_config.alg                                    # 选择的去噪算法（'origin'/ 'maskgit_plus'/ 'topk_margin'/ 'entropy'）
-        alg_temp = generation_config.alg_temp                          # 针对某些算法（margin/entropy）调整置信度的温度参数
-        temperature = generation_config.temperature                    # 采样时的温度
-        top_p = generation_config.top_p                                # top-p 截断采样参数
-        top_k = generation_config.top_k                                # top-k 截断采样参数
-
-        # histories 用于保存每一步的 x，如果需要返回历史则初始化为列表，否则为 None
+        output_history = True
+        return_dict_in_generate = generation_config.return_dict_in_generate  
+        max_length = generation_config.max_length                     
+        mask_token_id = generation_config.mask_token_id               
+        max_new_tokens = generation_config.max_new_tokens             
+        steps = min(generation_config.steps, max_new_tokens) 
+        eps = generation_config.eps                                    
+        alg = generation_config.alg if alg is None else alg                       
+        print("denoise algorithm: " + alg)
+        alg_temp = generation_config.alg_temp                          
+        temperature = generation_config.temperature                    
+        top_p = generation_config.top_p                                
+        top_k = generation_config.top_k              
+
         histories = [] if (return_dict_in_generate and output_history) else None
+        all_logit = []
+        generated_tokens = []
+        block_size = max_new_tokens if block_size < 0 else block_size
 
-        # 2. 如果没有传入 input_ids，而是直接传了 inputs_embeds，就根据 inputs_embeds 构造一个 placeholder 的 input_ids
         if input_ids is None:
             assert device is not None
             assert inputs_embeds is not None
-            bsz, seq_len = inputs_embeds.shape[:2]                   # batch size 和前缀长度
-            max_length = seq_len + max_new_tokens                     # 重新计算 max_length
-            # 创建一个全 0 的张量作为占位，后续会把 embedding 覆盖回去
+            bsz, seq_len = inputs_embeds.shape[:2]                  
+            max_length = seq_len + max_new_tokens                   
             input_ids = torch.full((bsz, seq_len), 0, dtype=torch.long).to(device)
 
-        # tok_idx 和 past_key_values 暂时留空，后面 prefix_lm 分支会用到
         tok_idx = None
         past_key_values = None
 
-        # 3. 把 input_ids pad 到 max_length，后面补 [MASK]
-        #    F.pad 的 (0, L) 表示在右侧 pad 长度为 (max_length - seq_len)，值为 mask_token_id
-        # import pdb;pdb.set_trace()
-        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id)  # 生成初始的 […, MASK, MASK, …]
-
-        # 4. 如果启用 prefix_lm 模式，先用 inputs_embeds 做一次常规模型前缀推理，得到 past_key_values 和首个 token
-        if prefix_lm:
-            dtype = inputs_embeds.dtype
-            # 先做一次前缀推理，use_cache=True 以获取 past_key_values
-            prefill = self.forward_dream(
-                None, attention_mask, tok_idx,
-                inputs_embeds=inputs_embeds.to(dtype),
-                use_cache=True
-            )
-            past_key_values = prefill.past_key_values
-            # 把前缀阶段模型最后一步的预测 token 取出，作为去噪的第一个位置
-            first_token = prefill.logits[:, -1:].argmax(dim=-1)  # 形状为 [B, 1]
-            # 只保留 mask 区域（原 x 的 right half）
-            x = x[:, input_ids.shape[1]:]                       # 形状 [B, max_new_tokens]
-            # 把 mask 区域第一位填为 first_token
-            x[:, :1] = first_token
-
-
-        #. prefill['logits'].shape.  torch.Size([1, 1063, 151667]) 即输入是这个
-
-        # 5. 当前不支持带 attention_mask 的情形，断言确保 attention_mask 一定为 None
-        assert attention_mask is None
+        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id) 
 
-        # 6. 构造去噪时刻表 timesteps，线性从 1 -> eps，共 (steps + 1) 个值
-        #    timesteps[i] 对应上一步噪声权重，timesteps[i+1] 对应本步噪声权重
         timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
-        # import pdb;pdb.set_trace()
-        # 7. 给用户一个机会在第 0 步“初始 x”阶段插入自定义逻辑
         x = generation_tokens_hook_func(None, x, None)
 
-        # 8. 如果用户指定 step_ratio，就根据比例重计算步数
         if step_ratio is not None:
             steps = int(max_new_tokens * step_ratio)
 
-        # 9. 计算每一步要去噪多少个 mask（如果传了 schedule，就用自定义调度）
-        if schedule is None:
-            sch = None
-        else:
-            # get_num_transfer_tokens_sch 返回形状 [B, steps] 的矩阵
-            sch = get_num_transfer_tokens_sch((x == mask_token_id), steps, schedule, schedule_kwargs)
-
-        # 10. 进入去噪主循环
-        for i in range(steps):
-            # 10.1 找出当前仍是 [MASK] 的位置，mask_index 为布尔矩阵 [B, current_length]
-            mask_index = (x == mask_token_id)
+        if add_boa_token:
+            bos_index = int((x.shape[1] - (x == mask_token_id).sum()) + (x == mask_token_id).sum() * 0.2)
+            x[:, bos_index] = 151684 # <|begin_of_audio|>
+
+        input_x = x.clone()
+        total_steps = steps
+        block_num = (x == mask_token_id).sum() // block_size
+        if block_num * block_size < (x == mask_token_id).sum(): block_num += 1
+        input_length = input_ids.shape[-1]
+
+        task = None
+        if "task" in kwargs: task = kwargs['task']
+        if cfg > 0:
+            import random
+            empty_prompt = ""
+            if task == "S2I":
+                empty_prompt = "<|im_start|>system\nPlease generate an image based on the input audio.<|im_end|>\n"
+                empty_prompt += "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n"
+                un_x = kwargs['tokenizer'].encode(empty_prompt)
+            elif task == "T2I":
+                empty_prompt = "<|im_start|>user\nGenerate an image based on the provided text description.\n"
+                empty_prompt += "<|im_end|>\n<|im_start|>assistant\n"
+                first_audio_token = kwargs['tokenizer'].encode("<|begin_of_audio|>")[0]
+                un_x_text = random.sample([_ for _ in range(first_audio_token)], 
+                                          input_ids.shape[1] - len(kwargs['tokenizer'].encode(empty_prompt)))
+                un_x = kwargs['tokenizer'].encode("<|im_start|>user\nGenerate an image based on the provided \
+                                                  text description.\n") 
+                un_x = un_x + un_x_text + kwargs['tokenizer'].encode("<|im_end|>\n<|im_start|>assistant\n")
+
+        for block_idx in range(block_num):
+            block_mask = torch.zeros([x.shape[-1]]).to(torch.bool).to(x.device)
+            block_mask[input_length + block_idx * block_size: input_length + (block_idx + 1) * block_size] = True
+            steps = int(block_mask.sum() / (x.shape[-1] - input_length) * total_steps)
+            timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
+            for i in tqdm(range(steps)):
+                mask_index = (x == mask_token_id)
+                if mask_index.sum() == 0: break
+                inputs_embeds_curr = self.model.embed_tokens(x)
 
-            # 10.2 先把 x 转成 embedding，得到形状 [B, current_length, D]
-            inputs_embeds_curr = self.model.embed_tokens(x)
-
-            # 10.3 如果非 prefix_lm，且外部传入了 inputs_embeds，则把前缀部分覆盖回去
-            if not prefix_lm:
                 if inputs_embeds is not None:
                     inputs_embeds_curr[:, :inputs_embeds.shape[1]] = inputs_embeds
 
-                # 用当前 embedding 做一次前向，得到 logits，形状 [B, current_length, V]
-                logits = self.forward_dream(None, attention_mask, tok_idx, inputs_embeds=inputs_embeds_curr).logits
-                # 把 logits 拼接成对齐当前预测：logits[:,1:] 对齐到 x[:, :-1]
-                logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
-            else:
-                # prefix_lm 模式，用 past_key_values 加速推理
-                logits = self.forward_dream(
-                    None, attention_mask, tok_idx,
-                    inputs_embeds=inputs_embeds_curr,
-                    past_key_values=past_key_values
-                ).logits
-                logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
-
-            # 10.4 用户自定义 logits 钩子，可以修改 logits 分布
-            # import pdb;pdb.set_trace()
-            logits = generation_logits_hook_func(i, x, logits)
-
-            # 10.5 取出当前所有 [MASK] 位置对应的 logits，形状 [num_mask, V]
-            mask_logits = logits[mask_index]
-
-            # 10.6 从 timesteps 中取出噪声权重 t, s
-            t = timesteps[i]
-            s = timesteps[i + 1]
-
-            # 10.7 根据不同算法决定本轮去噪逻辑
-            if alg == 'origin':
-                # 基础扩散算法：按概率 p_transfer 随机把一部分 mask 位置替换成 token
-                p_transfer = 1 - s / t if i < steps - 1 else 1  # 最后一轮保证把所有剩余 mask 都去掉
-                # x0 临时占位，全填 mask
-                x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id
-                # 随机采样哪些位置在本轮去噪：如果 torch.rand < p_transfer 就先去噪
-                transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer
-                # 对这些选中的位置，从 mask_logits 中采样真实 token
-                _, x0[transfer_index_t_s] = sample_tokens(
-                    mask_logits[transfer_index_t_s],
-                    temperature=temperature,
-                    top_p=top_p,
-                    top_k=top_k
-                )
-                # 更新 x：只替换 mask_index 位置
-                x[mask_index] = x0.clone()
-            else:
-                # MaskGIT+ / Top-K Margin / Entropy 算法
-                if alg == 'maskgit_plus':
-                    # 返回 confidence（置信度）和 x0（最可能的 token ID）
-                    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,
+                if cfg > 0:
+                    input_un_x = torch.tensor(un_x).unsqueeze(0).to(x.dtype).to(x.device)
+                    input_un_x = torch.cat([input_un_x, x[:, input_ids.shape[1]:]], dim=1)
+                    un_inpus_embeds = self.model.embed_tokens(input_un_x)
+
+                    attention_mask_cond = torch.ones([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
+                    attention_mask_cond = attention_mask_cond.to(torch.bool).to(inputs_embeds_curr.device)
+                    attention_mask_uncond = torch.zeros([1, inputs_embeds_curr.shape[1], inputs_embeds_curr.shape[1]])
+                    attention_mask_uncond[:, :un_inpus_embeds.shape[1], :un_inpus_embeds.shape[1]] = 1
+                    attention_mask_uncond = attention_mask_uncond.to(torch.bool).to(inputs_embeds.device)
+                    attention_mask = torch.cat([attention_mask_cond, attention_mask_uncond])
+                    attention_mask = attention_mask.unsqueeze(1)
+
+                    if inputs_embeds_curr.shape[1] != un_inpus_embeds.shape[1]:
+                        un_inpus_embeds = torch.cat([un_inpus_embeds, 
+                            torch.zeros_like(inputs_embeds_curr[:, :inputs_embeds_curr.shape[1] - 
+                                                                un_inpus_embeds.shape[1], :])], dim=1)
+                    input_inputs_embeds_curr = torch.cat([inputs_embeds_curr, un_inpus_embeds])
+
+                    model_logits = self.forward_dream(None, attention_mask, tok_idx, 
+                                                      inputs_embeds=input_inputs_embeds_curr).logits 
+                    logits = model_logits[:1]; un_logits = model_logits[1:]
+                    logits = un_logits + (cfg + 1) * (logits - un_logits)
+                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
+
+                else:
+                    logits = self.forward_dream(None, attention_mask, tok_idx, 
+                                                inputs_embeds=inputs_embeds_curr).logits
+                    logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
+
+                logits = generation_logits_hook_func(i, x, logits)
+
+                mask_logits = logits[mask_index]
+                if i == 0:
+                    input_index = torch.where(mask_index[0]==True)[0][0]
+
+                t = timesteps[i]
+                s = timesteps[i + 1]
+
+                if alg == 'origin':
+                    p_transfer = 1 - s / t if i < steps - 1 else 1  
+                    x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id
+                    transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer
+                    _, x0[transfer_index_t_s] = sample_tokens(
+                        mask_logits[transfer_index_t_s],
                         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
+                        max_position_penalty=max_position_penalty,
                     )
+                    x[mask_index] = x0.clone()
+
                 else:
-                    raise RuntimeError(f"Unknown alg: {alg}")
+                    if alg == 'maskgit_plus':
+                        confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, 
+                                                       max_position_penalty=max_position_penalty)
+                    elif alg == 'topk_margin':
+                        confidence, x0 = sample_tokens(
+                            mask_logits,
+                            temperature=temperature,
+                            top_p=top_p,
+                            top_k=top_k,
+                            margin_confidence=True,
+                            max_position_penalty=max_position_penalty,
+                        )
+                    elif alg == 'entropy':
+                        confidence, x0 = sample_tokens(
+                            mask_logits,
+                            temperature,
+                            top_p=top_p,
+                            top_k=top_k,
+                            neg_entropy=True,
+                            max_position_penalty=max_position_penalty,
+                        )
+                    elif alg == "entropy-penalty": 
+                        confidence, x0 = sample_tokens(
+                            mask_logits,
+                            temperature,
+                            top_p=top_p,
+                            top_k=top_k,
+                            neg_entropy=True,
+                            repeat_penalty=repeat_penalty if len(histories) != 0 else 1.0,
+                            past_x=histories[-1] if len(histories) != 0 else [],
+                            mask_id=mask_token_id,
+                            max_position_penalty=max_position_penalty,
+                        )
+                    else:
+                        raise RuntimeError(f"Unknown alg: {alg}")
 
-                # 当前还有多少 mask 位置
-                num_mask_token = mask_index.sum()
-                # 根据 schedule（或默认比例）决定本轮要去噪多少个
+                    block_mask_1 = block_mask[mask_index[0]]
+                    confidence = confidence + torch.where(block_mask_1, 0, -torch.inf).to(confidence.device)
 
-                if sch is not None:
-                    number_transfer_tokens = sch[0, i]
-                else:
+                    num_mask_token = mask_index.sum()
+                    num_mask_token = (x[:, block_mask] == mask_token_id).sum()
                     number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else num_mask_token
+                    number_transfer_tokens = max(number_transfer_tokens, 1)
 
-                if number_transfer_tokens > 0:
-                    if alg_temp is None or alg_temp == 0:
-                        # 直接选置信度最高的 number_transfer_tokens 个位置
-                        _, transfer_index = torch.topk(confidence, number_transfer_tokens)
-                    else:
-                        # 用温度调节 confidence，再按多项式采样 number_transfer_tokens 个
-                        confidence = confidence / alg_temp
-                        confidence = F.softmax(confidence, dim=-1)
-                        transfer_index = torch.multinomial(confidence, num_samples=number_transfer_tokens)
-
-                    # x0_ 临时占位，全填 mask
-                    x0_ = torch.zeros_like(x0, device=self.device, dtype=torch.long) + mask_token_id
-                    # 在选中的位置填入从 x0 (argmax token) 中取得的 token
-                    x0_[transfer_index] = x0[transfer_index].clone()
-
-
-
-                    # 更新 x：只替换 mask_index 位置
-                    x[mask_index] = x0_
-
-                    #如果出现的token有 151643（eos) ,那么他后面的所有都换成 151643，不需要再次mask
-                    SPECIAL_TOKEN_ID = 151643
-                    if (x == SPECIAL_TOKEN_ID).any():
-                        # 对每个 batch 处理
-                        for b in range(x.shape[0]):
-                            row = x[b]
-                            # 找到第一个出现 SPECIAL_TOKEN_ID 的位置
-                            idx = (row == SPECIAL_TOKEN_ID).nonzero(as_tuple=True)[0]
-                            if len(idx) > 0:
-                                first_idx = idx[0].item()
-                                # 该位置及其后面全部赋值为 SPECIAL_TOKEN_ID
-                                row[first_idx:] = SPECIAL_TOKEN_ID
-                                x[b] = row
-
-            # 10.8 用户自定义 token 钩子：对本轮更新后的 x 做额外处理
-            x = generation_tokens_hook_func(i, x, logits)
-
-            # 10.9 如果需要保存历史，就把当前 x clone 一份放进去
-            if histories is not None:
-                histories.append(x.clone())
-
-
-        # ForkedPdb().set_trace()   
-        # 11. 循环结束后，根据 return_dict_in_generate 决定返回形式
-        if return_dict_in_generate:
-            return DreamModelOutput(
-                sequences=x,     # 最终生成的完整 token 序列 [B, max_length]
-                history=histories,  # 如果启用，会包含每一步的 x
-            )
-        else:
-            return x  # 只返回最终序列 [B, max_length]
+                    if number_transfer_tokens > 0:
+                        if alg_temp is None or alg_temp == 0:
+                            _, transfer_index = torch.topk(confidence, number_transfer_tokens)
+                        else:
+                            confidence = confidence / alg_temp
+                            confidence = F.softmax(confidence, dim=-1)
+                            transfer_index = torch.multinomial(confidence, num_samples=number_transfer_tokens)
+
+                        x0_ = torch.zeros_like(x0, device=self.device, dtype=torch.long) + mask_token_id
+                        x0_[transfer_index] = x0[transfer_index].clone()
+                        x[mask_index] = x0_
+
+                        logit,indic = torch.max(torch.softmax(logits.clone(),dim=-1),-1)
+                        logit = logit[0][x[0]!=0]
+                        indic = indic[0][x[0]!=0]
+                        temp_X = x[0][x[0]!=0]
+
+                x = generation_tokens_hook_func(i, x, logits)
+
+                if histories is not None:
+                    histories.append(x.clone())
+                    all_logit.append(torch.max(logits.clone(),-1)[-1])
+
+        return (x, histories)