From caab8f746b397028798795ee73a4e0f0a1f71e54 Mon Sep 17 00:00:00 2001
From: misyaguziya <misyaguziya@gmail.com>
Date: Wed, 28 Aug 2024 11:39:34 +0900
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=8D=EF=B8=8F[Update]=20Controller=20:?=
 =?UTF-8?q?=20=E8=B5=B7=E5=8B=95=E6=99=82=E3=81=AB=E3=83=A2=E3=83=87?=
 =?UTF-8?q?=E3=83=AB=E3=81=AEweight=E3=83=95=E3=82=A1=E3=82=A4=E3=83=AB?=
 =?UTF-8?q?=E3=82=92=E3=83=80=E3=82=A6=E3=83=B3=E3=83=AD=E3=83=BC=E3=83=89?=
 =?UTF-8?q?=E3=81=99=E3=82=8B=E3=82=88=E3=81=86=E3=81=AB=E5=AE=9F=E8=A3=85?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src-python/webui_controller.py | 28 +++++++++++++++++++++++++++-
 src-python/webui_mainloop.py   | 14 ++++++++------
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/src-python/webui_controller.py b/src-python/webui_controller.py
index 47003ba5..515e4e7d 100644
--- a/src-python/webui_controller.py
+++ b/src-python/webui_controller.py
@@ -1292,7 +1292,7 @@ def getListInputDevice(*args, **kwargs) -> dict:
 def getListOutputDevice(*args, **kwargs) -> dict:
     return {"status":200, "result": model.getListOutputDevice()}
 
-def init():
+def init(endpoints:dict, *args, **kwargs) -> None:
     print(json.dumps({"status":348, "log": "Start Initialization"}), flush=True)
     print(json.dumps({"status":348, "log": "Start InitSetTranslateEngine"}), flush=True)
     initSetTranslateEngine()
@@ -1311,6 +1311,19 @@ def init():
     print(json.dumps({"status":348, "log": "Set Translation Engine"}), flush=True)
     updateTranslationEngineAndEngineList()
 
+    # check Downloaded CTranslate2 Model Weight
+    print(json.dumps({"status":348, "log": "Check Downloaded CTranslate2 Model Weight"}), flush=True)
+    if config.USE_TRANSLATION_FEATURE is True and model.checkCTranslatorCTranslate2ModelWeight() is False:
+        def callback(progress):
+            print(json.dumps({
+                "endpoint":endpoints["ctranslate2"],
+                "status":200,
+                "result":{
+                    "progress":progress
+                    }
+            }), flush=True)
+        model.downloadCTranslate2ModelWeight(callback)
+
     # set Transcription Engine
     print(json.dumps({"status":348, "log": "Set Transcription Engine"}), flush=True)
     if config.USE_WHISPER_FEATURE is True:
@@ -1318,6 +1331,19 @@ def init():
     else:
         config.SELECTED_TRANSCRIPTION_ENGINE = "Google"
 
+    # check Downloaded Whisper Model Weight
+    print(json.dumps({"status":348, "log": "Check Downloaded Whisper Model Weight"}), flush=True)
+    if config.USE_WHISPER_FEATURE is True and model.checkTranscriptionWhisperModelWeight() is False:
+        def callback(progress):
+            print(json.dumps({
+                "endpoint":endpoints["whisper"],
+                "status":200,
+                "result":{
+                    "progress":progress
+                    }
+            }), flush=True)
+        model.downloadWhisperModelWeight(callback)
+
     # set word filter
     print(json.dumps({"status":348, "log": "Set Word Filter"}), flush=True)
     model.addKeywords()
diff --git a/src-python/webui_mainloop.py b/src-python/webui_mainloop.py
index 485e5057..a6d73a19 100644
--- a/src-python/webui_mainloop.py
+++ b/src-python/webui_mainloop.py
@@ -326,14 +326,18 @@ def main():
         print(response, flush=True)
 
 if __name__ == "__main__":
+    print(json.dumps({"status":200, "endpoint": "/initialization/start", "result":True}), flush=True)
+    controller.init({
+        "ctranslate2": action_mapping["/controller/callback_download_ctranslate2_weight"]["download"],
+        "whisper": action_mapping["/controller/callback_download_whisper_weight"]["download"],
+    })
+    print(json.dumps({"status":200, "endpoint": "/initialization/completed", "result":True}), flush=True)
+    print(json.dumps({"status":348, "log": "Initialization from Python."}), flush=True)
+
     process = "main"
     match process:
         case "main":
             try:
-                print(json.dumps({"status":200, "endpoint": "/initialization/start", "result":True}), flush=True)
-                controller.init()
-                print(json.dumps({"status":200, "endpoint": "/initialization/completed", "result":True}), flush=True)
-                print(json.dumps({"status":348, "log": "Initialization from Python."}), flush=True)
                 while True:
                     main()
             except Exception:
@@ -342,7 +346,6 @@ if __name__ == "__main__":
                     traceback.print_exc(file=f)
 
         case "test":
-            controller.init()
             response_data, status = handleControllerRequest("/controller/callback_download_ctranslate2_weight")
             response = {
                 "status": status,
@@ -360,7 +363,6 @@ if __name__ == "__main__":
 
         case "test_all":
             import time
-            controller.init()
             for endpoint, value in config_mapping.items():
                 response_data, status = handleConfigRequest(endpoint)
                 response = {