Fix batch generation

generation_utils.py

@@ -433,18 +433,21 @@ class DreamGenerationMixin:
                     confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
                 else:
                     raise RuntimeError(f"Unknown alg: {alg}")
-                num_mask_token = mask_index.sum()
-                number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else num_mask_token
+                num_mask_token = mask_index.sum() / mask_index.shape[0]
+                number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)
+                full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
+                full_confidence[mask_index] = confidence
                 if number_transfer_tokens > 0:
                     if alg_temp is None or alg_temp == 0:
-                        _, transfer_index = torch.topk(confidence, number_transfer_tokens)
+                        _, transfer_index = torch.topk(full_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_
+                        transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
+                    x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
+                    x_[mask_index] = x0.clone()
+                    row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)
+                    x[row_indices,transfer_index] = x_[row_indices,transfer_index]
 
             # this allows user-defined token control of the intermediate steps
             x = generation_tokens_hook_func(i, x, logits)