From 83f2c1a8edd2f79af85e597a8789a0697f25b0bc Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Tue, 12 Dec 2023 22:11:23 +0200
Subject: [PATCH] #1524 Add FPS limiter to streamed rendering

---
 public/index.html            | 13 +++++++++++++
 public/script.js             |  8 +++++++-
 public/scripts/power-user.js | 12 +++++++++++-
 public/scripts/utils.js      | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/public/index.html b/public/index.html
index 7eeccf4e0..348a47fa3 100644
--- a/public/index.html
+++ b/public/index.html
@@ -2890,6 +2890,19 @@
                                         </div>
                                     </div>
                                 </div>
+                                <div id="streaming-fps" class="range-block">
+                                    <div class="range-block-title" data-i18n="Streaming FPS">
+                                        Streaming FPS
+                                    </div>
+                                    <div class="range-block-range-and-counter">
+                                        <div class="range-block-range">
+                                            <input type="range" id="streaming_fps" name="streaming_fps" min="5" max="100" step="5">
+                                        </div>
+                                        <div class="range-block-counter">
+                                            <input type="number" min="5" max="100" step="1" data-for="streaming_fps" id="streaming_fps_counter">
+                                        </div>
+                                    </div>
+                                </div>
                             </div>
                         </div>
                     </div>
diff --git a/public/script.js b/public/script.js
index bc29aff53..fdd5a76fd 100644
--- a/public/script.js
+++ b/public/script.js
@@ -143,6 +143,7 @@ import {
     onlyUnique,
     getBase64Async,
     humanFileSize,
+    Stopwatch,
 } from './scripts/utils.js';
 
 import { ModuleWorkerWrapper, doDailyExtensionUpdatesCheck, extension_settings, getContext, loadExtensionSettings, processExtensionHelpers, registerExtensionHelper, renderExtensionTemplate, runGenerationInterceptors, saveMetadataDebounced } from './scripts/extensions.js';
@@ -2805,7 +2806,10 @@ class StreamingProcessor {
         }
 
         try {
+            const sw = new Stopwatch(1000 / power_user.streaming_fps);
+            const timestamps = [];
             for await (const { text, swipes } of this.generator()) {
+                timestamps.push(Date.now());
                 if (this.isStopped) {
                     this.onStopStreaming();
                     return;
@@ -2813,8 +2817,10 @@ class StreamingProcessor {
 
                 this.result = text;
                 this.swipes = swipes;
-                this.onProgressStreaming(this.messageId, message_already_generated + text);
+                await sw.tick(() => this.onProgressStreaming(this.messageId, message_already_generated + text));
             }
+            const seconds = (timestamps[timestamps.length - 1] - timestamps[0]) / 1000;
+            console.warn(`Stream stats: ${timestamps.length} tokens, ${seconds.toFixed(2)} seconds, rate: ${Number(timestamps.length / seconds).toFixed(2)} TPS`);
         }
         catch (err) {
             console.error(err);
diff --git a/public/scripts/power-user.js b/public/scripts/power-user.js
index 9689fb212..2983047d9 100644
--- a/public/scripts/power-user.js
+++ b/public/scripts/power-user.js
@@ -114,6 +114,7 @@ let power_user = {
     },
     markdown_escape_strings: '',
     chat_truncation: 100,
+    streaming_fps: 30,
 
     ui_mode: ui_mode.POWER,
     fast_ui_mode: true,
@@ -1460,6 +1461,9 @@ function loadPowerUserSettings(settings, data) {
     $('#chat_truncation').val(power_user.chat_truncation);
     $('#chat_truncation_counter').val(power_user.chat_truncation);
 
+    $('#streaming_fps').val(power_user.streaming_fps);
+    $('#streaming_fps_counter').val(power_user.streaming_fps);
+
     $('#font_scale').val(power_user.font_scale);
     $('#font_scale_counter').val(power_user.font_scale);
 
@@ -2701,6 +2705,12 @@ $(document).ready(() => {
         saveSettingsDebounced();
     });
 
+    $('#streaming_fps').on('input', function () {
+        power_user.streaming_fps = Number($('#streaming_fps').val());
+        $('#streaming_fps_counter').val(power_user.streaming_fps);
+        saveSettingsDebounced();
+    });
+
     $('input[name="font_scale"]').on('input', async function (e) {
         power_user.font_scale = Number(e.target.value);
         $('#font_scale_counter').val(power_user.font_scale);
@@ -3134,7 +3144,7 @@ $(document).ready(() => {
         saveSettingsDebounced();
     });
 
-    $('#reduced_motion').on('input', function() {
+    $('#reduced_motion').on('input', function () {
         power_user.reduced_motion = !!$(this).prop('checked');
         localStorage.setItem(storage_keys.reduced_motion, String(power_user.reduced_motion));
         switchReducedMotion();
diff --git a/public/scripts/utils.js b/public/scripts/utils.js
index 6ea0508d3..dfaba74ce 100644
--- a/public/scripts/utils.js
+++ b/public/scripts/utils.js
@@ -741,6 +741,38 @@ export function escapeRegex(string) {
     return string.replace(/[/\-\\^$*+?.()|[\]{}]/g, '\\$&');
 }
 
+export class Stopwatch {
+    /**
+     * Initializes a Stopwatch class.
+     * @param {number} interval Update interval in milliseconds. Must be a finite number above zero.
+     */
+    constructor(interval) {
+        if (isNaN(interval) || !isFinite(interval) || interval <= 0) {
+            console.warn('Invalid interval for Stopwatch, setting to 1');
+            interval = 1;
+        }
+
+        this.interval = interval;
+        this.lastAction = Date.now();
+    }
+
+    /**
+     * Executes a function if the interval passed.
+     * @param {(arg0: any) => any} action Action function
+     * @returns Promise<void>
+     */
+    async tick(action) {
+        const passed = (Date.now() - this.lastAction);
+
+        if (passed < this.interval) {
+            return;
+        }
+
+        await action();
+        this.lastAction = Date.now();
+    }
+}
+
 /**
  * Provides an interface for rate limiting function calls.
  */