Just Tuple of Numpy

app.py

@@ -11,6 +11,7 @@ import torch
 import xxhash
 from datasets import Audio
 from transformers import AutoModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
 import io
 
 if gr.NO_RELOAD:
@@ -74,15 +75,16 @@ def response(state: AppState, audio: tuple):
     )
     if spaces.config.Config.zero_gpu:
         if state.model_outs is not None:
-            state.model_outs.past_key_values = tuple(
+            state.model_outs = tuple(
                 tuple(torch.tensor(vec).cuda() for vec in tup)
-                for tup in state.model_outs.past_key_values
+                for tup in state.model_outs
             )
-    prev_outs = state.model_outs
+    causal_outs = CausalLMOutputWithPast(past_key_values=state.model_outs)
+    prev_outs = causal_outs
     start = False
     for resp, outs in diva_audio(
         (state.sampling_rate, state.stream),
-        prev_outs=(prev_outs if state.model_outs is not None else None),
+        prev_outs=(prev_outs if prev_outs is not None else None),
     ):
         if not start:
             state.conversation.append({"role": "assistant", "content": resp})
@@ -94,10 +96,9 @@ def response(state: AppState, audio: tuple):
     del outs.logits
     del outs.hidden_states
     if spaces.config.Config.zero_gpu:
-        outs.past_key_values = tuple(
+        outs = tuple(
             tuple(vec.cpu().numpy() for vec in tup) for tup in outs.past_key_values
         )
-        print(outs)
     return (
         AppState(conversation=state.conversation, model_outs=outs),
         state.conversation,