From b6c407f51defa1d6fcb2341f31347b8bf8ab47e8 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 31 May 2023 23:41:42 -0300
Subject: [PATCH] Don't stream at more than 24 fps

This is a performance optimization
---
 modules/text_generation.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/modules/text_generation.py b/modules/text_generation.py
index 71decb0c..f4faf4cc 100644
--- a/modules/text_generation.py
+++ b/modules/text_generation.py
@@ -188,7 +188,19 @@ def _generate_reply(question, state, eos_token=None, stopping_strings=None, is_c
     shared.stop_everything = False
     clear_torch_cache()
     seed = set_manual_seed(state['seed'])
+    is_stream = state['stream']
+    last_update = -1
+    reply = ''
     for reply in generate_func(question, original_question, seed, state, eos_token, stopping_strings, is_chat=is_chat):
+        if is_stream:
+            cur_time = time.time()
+            if cur_time - last_update > 0.041666666666666664:  # Limit streaming to 24 fps
+                last_update = cur_time
+                yield reply
+        else:
+            yield reply
+
+    if is_stream:
         yield reply