Update app.py

app.py

@@ -93,34 +93,57 @@ class VoiceMapper:
         return default_voice
 
 
-# Monkey patch the _update_model_kwargs_for_generation function to handle dict outputs
-def _update_model_kwargs_for_generation_fixed(
+# Patch the _update_model_kwargs_for_generation method
+def patched_update_model_kwargs_for_generation(
+    self,
     outputs,
     model_kwargs,
     is_encoder_decoder=False,
-    **kwargs,
+    model_inputs=None,
+    num_new_tokens=1,
 ):
-    """Fixed version that handles both dict and object-like outputs"""
-    # Update past_key_values - handle both dict and object-like outputs
+    """Patched version that handles both dict and object-like outputs"""
+    # Handle both dict and object-like outputs for cache
+    cache_name = "past_key_values"
+    
     if isinstance(outputs, dict):
-        model_kwargs["past_key_values"] = outputs.get("past_key_values")
+        # For dict outputs, use .get() method
+        model_kwargs[cache_name] = outputs.get(cache_name)
     else:
-        model_kwargs["past_key_values"] = getattr(outputs, "past_key_values", None)
-
-    # Update attention mask
-    if "attention_mask" in model_kwargs:
-        attention_mask = model_kwargs["attention_mask"]
-        model_kwargs["attention_mask"] = torch.cat(
-            [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))],
-            dim=-1,
-        )
+        # For object outputs, try to get the attribute
+        model_kwargs[cache_name] = getattr(outputs, cache_name, None)
+
+    if getattr(self, "config", None) is not None:
+        if "token_type_ids" in model_kwargs and model_kwargs["token_type_ids"] is not None:
+            token_type_ids = model_kwargs["token_type_ids"]
+            model_kwargs["token_type_ids"] = torch.cat(
+                [token_type_ids, token_type_ids[:, -1:]], dim=-1
+            )
 
-    return model_kwargs
+        if not is_encoder_decoder:
+            # update attention mask
+            if "attention_mask" in model_kwargs and model_kwargs["attention_mask"] is not None:
+                attention_mask = model_kwargs["attention_mask"]
+                model_kwargs["attention_mask"] = torch.cat(
+                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))],
+                    dim=-1,
+                )
+        else:
+            # update decoder attention mask
+            if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None:
+                decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+                model_kwargs["decoder_attention_mask"] = torch.cat(
+                    [
+                        decoder_attention_mask,
+                        decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1)),
+                    ],
+                    dim=-1,
+                )
 
+    if model_inputs is not None and "cache_position" in model_inputs:
+        model_kwargs["cache_position"] = model_inputs["cache_position"][-1:] + num_new_tokens
 
-# Apply the monkey patch
-import vibevoice.modular.modeling_vibevoice_streaming_inference as inference_module
-inference_module._update_model_kwargs_for_generation = _update_model_kwargs_for_generation_fixed
+    return model_kwargs
 
 
 # Check if CUDA is available
@@ -147,6 +170,9 @@ MODEL = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
     attn_implementation="sdpa",
 )
 
+# Apply the patch to the model instance
+MODEL._update_model_kwargs_for_generation = patched_update_model_kwargs_for_generation.__get__(MODEL, type(MODEL))
+
 MODEL.eval()
 MODEL.set_ddpm_inference_steps(num_steps=5)