Upload folder using huggingface_hub

generation_utils.py

@@ -425,7 +425,13 @@ class DreamGenerationMixin:
             logits = generation_logits_hook_func(i, x, logits)
             probs = F.softmax(logits, dim=-1)
                         
-
+            #ВОТ ТУТ МАСКИНГ ДОБАВЛЯЕМ!
+            keep_mask = torch.ones_like(x, device=self.device, dtype=torch.bool)
+            if remask_schedule and remask_schedule == "tau":
+                tau = remask_tau[i]
+                token_confidence = probs.amax(-1)
+                keep_mask = token_confidence >= tau
+                x = torch.where(keep_mask, x, mask_token_id)
 
 
             mask_logits = logits[mask_index]
@@ -463,13 +469,6 @@ class DreamGenerationMixin:
                     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]
 
-            #ВОТ ТУТ МАСКИНГ ДОБАВЛЯЕМ!
-            keep_mask = torch.ones_like(x, device=self.device, dtype=torch.bool)
-            if remask_schedule and remask_schedule == "tau":
-                tau = remask_tau[i]
-                token_confidence = probs.amax(-1)
-                keep_mask = token_confidence >= tau
-                x = torch.where(keep_mask, x, mask_token_id)
             # this allows user-defined token control of the intermediate steps
             x = generation_tokens_hook_func(i, x, logits)