From 9777bf99ad2d12e2b536e9e69e8fcbf5d45b3f21 Mon Sep 17 00:00:00 2001
From: marijnvg-tng <marijn.vangeest@tngtech.com>
Date: Wed, 20 Nov 2024 15:34:26 +0100
Subject: [PATCH] 26 - Code cleanup after revision

---
 WebUI/src/App.vue                       |  6 ++++++
 WebUI/src/assets/js/store/models.ts     |  1 -
 WebUI/src/components/AddLLMDialog.vue   | 10 ++++++----
 WebUI/src/components/DownloadDialog.vue |  4 ----
 WebUI/src/components/InpaintMask.vue    |  1 -
 WebUI/src/components/WarningDialog.vue  |  9 +++++----
 WebUI/src/views/Answer.vue              | 12 ------------
 service/model_downloader.py             | 14 +++++++++++---
 service/web_api.py                      |  7 ++-----
 9 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/WebUI/src/App.vue b/WebUI/src/App.vue
index d025f238..7a5c7087 100644
--- a/WebUI/src/App.vue
+++ b/WebUI/src/App.vue
@@ -239,6 +239,9 @@ function showDownloadModelConfirm(downList: DownloadModelParam[], success?: () =
 
 function showModelRequest() {
   showModelRequestDialog.value = true;
+  nextTick(() => {
+    addLLMCompt.value!.onShow();
+  });
 }
 
 function callCheckModel(){
@@ -249,6 +252,9 @@ function showWarning(message : string, func : () => void) {
   warningCompt.value!.warningMessage = message;
   showWarningDialog.value = true;
   warningCompt.value!.confirmFunction = func;
+  nextTick(() => {
+    warningCompt.value!.onShow();
+  });
 }
 
 </script>
diff --git a/WebUI/src/assets/js/store/models.ts b/WebUI/src/assets/js/store/models.ts
index 62db3420..4f96b650 100644
--- a/WebUI/src/assets/js/store/models.ts
+++ b/WebUI/src/assets/js/store/models.ts
@@ -46,7 +46,6 @@ export const useModels = defineStore("models", () => {
         const notYetDownloaded = (model: Model) => !downloadedModels.map(m => m.name).includes(model.name);
 
         models.value = [...downloadedModels, ...userModels, ...predefinedModels.filter(notYetDownloaded)];
-        console.log(models);
 
         }
 
diff --git a/WebUI/src/components/AddLLMDialog.vue b/WebUI/src/components/AddLLMDialog.vue
index e829530d..8a1d6638 100644
--- a/WebUI/src/components/AddLLMDialog.vue
+++ b/WebUI/src/components/AddLLMDialog.vue
@@ -34,9 +34,6 @@ const emits = defineEmits<{
     (e: "showWarning", warning : string, func : () => void): void
 }>();
 
-onDeactivated(() => {
-    animate.value = false;
-})
 
 function fastGenerate(e: KeyboardEvent) {
   if (e.code == "Enter") {
@@ -49,6 +46,10 @@ function fastGenerate(e: KeyboardEvent) {
   }
 }
 
+function onShow() {
+  animate.value = true;
+}
+
 async function addModel() {
   const previousModel = globalSetup.modelSettings.llm_model
   const is_in_models = models.models.some((model) => model.name === modelRequest.value)
@@ -67,7 +68,6 @@ async function addModel() {
         globalSetup.modelSettings.llm_model = previousModel
         addModelErrorMessage.value = i18nState.ERROR_REPO_NOT_EXISTS
         addModelError.value = true;
-        console.log("Hey")
       }
   } else {
     globalSetup.modelSettings.llm_model = previousModel
@@ -112,6 +112,8 @@ function closeAdd() {
   emits("close");
 }
 
+defineExpose({ onShow });
+
 </script>
 <style scoped>
 table {
diff --git a/WebUI/src/components/DownloadDialog.vue b/WebUI/src/components/DownloadDialog.vue
index 7488bcf5..81a9c3ce 100644
--- a/WebUI/src/components/DownloadDialog.vue
+++ b/WebUI/src/components/DownloadDialog.vue
@@ -111,10 +111,6 @@ const readTerms = ref(false);
 const downloadList = ref<DownloadModelRender[]>([]);
 
 
-onDeactivated(() => {
-    animate.value = false;
-})
-
 function dataProcess(line: string) {
     console.log(line);
     const dataJson = line.slice(5);
diff --git a/WebUI/src/components/InpaintMask.vue b/WebUI/src/components/InpaintMask.vue
index c660d3da..d8b07ec5 100644
--- a/WebUI/src/components/InpaintMask.vue
+++ b/WebUI/src/components/InpaintMask.vue
@@ -108,7 +108,6 @@ function moveWithShadow(e: MouseEvent) {
 
 function showShadow() {
     shadowPos.show = true;
-    console.log("showShadow");
 }
 
 defineExpose({ clearMaskImage, getMaskImage })
diff --git a/WebUI/src/components/WarningDialog.vue b/WebUI/src/components/WarningDialog.vue
index 54a6a533..86ce3b97 100644
--- a/WebUI/src/components/WarningDialog.vue
+++ b/WebUI/src/components/WarningDialog.vue
@@ -22,9 +22,6 @@ const emits = defineEmits<{
     (e: "close"): void
 }>();
 
-onDeactivated(() => {
-    animate.value = false;
-})
 
 async function confirmAdd() {
   confirmFunction.value()
@@ -35,7 +32,11 @@ function cancelConfirm() {
   emits("close");
 }
 
-defineExpose({warningMessage, confirmFunction});
+function onShow(){
+  animate.value = true
+}
+
+defineExpose({warningMessage, confirmFunction, onShow });
 
 </script>
 <style scoped>
diff --git a/WebUI/src/views/Answer.vue b/WebUI/src/views/Answer.vue
index 1b05afa6..603ff59c 100644
--- a/WebUI/src/views/Answer.vue
+++ b/WebUI/src/views/Answer.vue
@@ -176,12 +176,6 @@
             </div>
             <rag v-if="ragData.showUploader" ref="ragPanel" @close="ragData.showUploader = false"></rag>
         </div>
-        <teleport to="#answerPanel" v-if="showDowloadDlg">
-            <download-dialog ref="downloadDigCompt"></download-dialog>
-        </teleport>
-        <teleport to="#answerPanel" v-if="showModelRequestDialog">
-            <add-l-l-m-dialog ref="addLLMDialog" @add-model="checkModel"></add-l-l-m-dialog>
-        </teleport>
     </div>
 
 </template>
@@ -190,8 +184,6 @@ import Rag from "../components/Rag.vue";
 import ProgressBar from "../components/ProgressBar.vue";
 import LoadingBar from "../components/LoadingBar.vue";
 import DropSelector from "@/components/DropSelector.vue";
-import DownloadDialog from "@/components/DownloadDialog.vue";
-import AddLLMDialog from "@/components/AddLLMDialog.vue";
 import ModelDropDownItem from "@/components/ModelDropDownItem.vue";
 import { useI18N } from '@/assets/js/store/i18n';
 import { toast } from '@/assets/js/toast';
@@ -226,15 +218,11 @@ const loadingModel = ref(false);
 let receiveOut = "";
 let chatPanel: HTMLElement;
 const markdownParser = new MarkdownParser(i18nState.COM_COPY);
-const showDowloadDlg = ref(false);
-const showModelRequestDialog = ref(false);
 const ragData = reactive({
     enable: false,
     processEnable: false,
     showUploader: false,
 });
-const downloadDigCompt = ref<InstanceType<typeof DownloadDialog>>()
-const addLLMDialog = ref<InstanceType<typeof AddLLMDialog>>()
 
 const source = ref("");
 const emits = defineEmits<{
diff --git a/service/model_downloader.py b/service/model_downloader.py
index 3d9a2f25..16d0ffc9 100644
--- a/service/model_downloader.py
+++ b/service/model_downloader.py
@@ -1,4 +1,5 @@
 from huggingface_hub import HfFileSystem, hf_hub_url, model_info
+from huggingface_hub.utils._errors import RepositoryNotFoundError
 from typing import Any, Callable, Dict, List
 from os import path, makedirs, rename
 import requests
@@ -11,11 +12,11 @@
 import traceback
 import concurrent.futures
 import aipg_utils as utils
+import shutil
 
 model_list_cache = dict()
 model_lock = Lock()
 
-
 class HFFileItem:
     relpath: str
     size: int
@@ -83,8 +84,12 @@ def __init__(self, hf_token=None) -> None:
         self.thread_lock = Lock()
         self.hf_token = hf_token
 
-    def probe_url(self, repo_id: str):
-        model_info(repo_id)
+    def hf_url_exists(self, repo_id: str):
+        try:
+            model_info(repo_id)
+            return True
+        except RepositoryNotFoundError:
+            return False
 
     def probe_type(self, repo_id : str):
         return model_info(repo_id).pipeline_tag
@@ -264,6 +269,9 @@ def multiple_thread_downlod(self, thread_count: int):
                     path.join(self.save_path, self.repo_id.replace("/", "---"))
                 ),
             )
+        else:
+            # Download aborted
+            shutil.rmtree(self.save_path_tmp)
 
     def start_report_download_progress(self):
         thread = Thread(target=self.report_download_progress)
diff --git a/service/web_api.py b/service/web_api.py
index e18182b3..0146baa1 100644
--- a/service/web_api.py
+++ b/service/web_api.py
@@ -195,15 +195,12 @@ def check_model_exist():
         result_list.append({"repo_id": repo_id, "type": type, "exist": exist})
     return jsonify({"code": 0, "message": "success", "exists": result_list})
 
+
 @app.route("/api/checkURLExists", methods=["POST"])
 def check_url_exists():
     address = request.get_json()
     downloader = HFPlaygroundDownloader()
-    exists = True
-    try:
-        downloader.probe_url(address)
-    except hf.utils._errors.RepositoryNotFoundError:
-        exists = False
+    exists = downloader.hf_url_exists(address)
 
     return jsonify(
             {