Upload model

generation_utils.py

@@ -302,8 +302,9 @@ class DreamGenerationMixin:
         **kwargs,
     ) -> Union[DreamModelOutput, torch.LongTensor]:
         # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
-        generation_hook_func = kwargs.pop("generation_hook_func", lambda x, logits: (x, logits))
         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
@@ -355,7 +356,8 @@ class DreamGenerationMixin:
             input_ids,
             attention_mask=attention_mask,
             generation_config=generation_config,
-            generation_hook_func=generation_hook_func
+            generation_tokens_hook_func=generation_tokens_hook_func,
+            generation_logits_hook_func=generation_logits_hook_func
         )
         return result
 
@@ -364,7 +366,8 @@ class DreamGenerationMixin:
         input_ids: torch.LongTensor,
         attention_mask: Optional[torch.LongTensor],
         generation_config: DreamGenerationConfig,
-        generation_hook_func
+        generation_tokens_hook_func,
+        generation_logits_hook_func
     ) -> Union[DreamModelOutput, torch.LongTensor]:
         # init values
         output_history = generation_config.output_history
@@ -400,10 +403,15 @@ class DreamGenerationMixin:
             attention_mask = "full"
 
         timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
+        
         for i in range(steps):
             mask_index = (x == mask_token_id)
             logits = self(x, attention_mask, tok_idx).logits
             logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)
+
+            # this allows user-defined logits control of the intermediate steps
+            logits = generation_logits_hook_func(i, x, logits)
+
             mask_logits = logits[mask_index]
             t = timesteps[i]
             s = timesteps[i + 1]
@@ -435,9 +443,9 @@ class DreamGenerationMixin:
                     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_
-                    
-            # this allows user-defined control of the intermediate steps
-            x, logits = generation_hook_func(x, logits)
+
+            # this allows user-defined token control of the intermediate steps
+            x = generation_tokens_hook_func(i, x, logits)
 
             if histories is not None:
                 histories.append(x.clone())