From 70167c0f4566e9ad1668cec4a02c3e3cc3d67ef9 Mon Sep 17 00:00:00 2001
From: James <namnh0122@gmail.com>
Date: Tue, 26 Nov 2024 10:52:21 +0700
Subject: [PATCH] fix: (#4091) model size not displaying

---
 engine/cli/commands/model_import_cmd.cc |  8 +-------
 engine/common/download_task.h           |  4 ++--
 engine/services/download_service.cc     | 16 +++++++++++-----
 engine/services/model_service.cc        |  6 ++++++
 engine/services/model_service.h         |  1 -
 5 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/engine/cli/commands/model_import_cmd.cc b/engine/cli/commands/model_import_cmd.cc
index 6abaad2c8..f8cf6a810 100644
--- a/engine/cli/commands/model_import_cmd.cc
+++ b/engine/cli/commands/model_import_cmd.cc
@@ -1,13 +1,7 @@
 #include "model_import_cmd.h"
-#include <filesystem>
-#include <vector>
-#include "config/gguf_parser.h"
-#include "config/yaml_config.h"
-#include "database/models.h"
+#include <json/value.h>
 #include "httplib.h"
-#include "json/json.h"
 #include "server_start_cmd.h"
-#include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
 
 namespace commands {
diff --git a/engine/common/download_task.h b/engine/common/download_task.h
index 94fb11a48..95e736394 100644
--- a/engine/common/download_task.h
+++ b/engine/common/download_task.h
@@ -2,9 +2,9 @@
 
 #include <json/json.h>
 #include <filesystem>
+#include <optional>
 #include <sstream>
 #include <string>
-#include <optional>
 
 enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };
 
@@ -161,4 +161,4 @@ inline DownloadTask GetDownloadTaskFromJson(const Json::Value item_json) {
   }
   return task;
 }
-}  // namespace common
\ No newline at end of file
+}  // namespace common
diff --git a/engine/services/download_service.cc b/engine/services/download_service.cc
index 09c4d1a75..d855c8f61 100644
--- a/engine/services/download_service.cc
+++ b/engine/services/download_service.cc
@@ -544,10 +544,16 @@ void DownloadService::EmitTaskCompleted(const std::string& task_id) {
 }
 
 void DownloadService::ExecuteCallback(const DownloadTask& task) {
-  std::lock_guard<std::mutex> lock(callbacks_mutex_);
-  auto it = callbacks_.find(task.id);
-  if (it != callbacks_.end()) {
-    it->second(task);
-    callbacks_.erase(it);
+  std::lock_guard<std::mutex> active_task_lock(active_tasks_mutex_);
+  if (auto it = active_tasks_.find(task.id); it != active_tasks_.end()) {
+    for (auto& item : it->second->items) {
+      item.downloadedBytes = item.bytes;
+    }
+    std::lock_guard<std::mutex> lock(callbacks_mutex_);
+    auto callback = callbacks_.find(task.id);
+    if (callback != callbacks_.end()) {
+      callback->second(*it->second);
+      callbacks_.erase(callback);
+    }
   }
 }
diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc
index 80d02f4c7..d6e66a717 100644
--- a/engine/services/model_service.cc
+++ b/engine/services/model_service.cc
@@ -463,6 +463,12 @@ ModelService::DownloadModelFromCortexsoAsync(
     yaml_handler.ModelConfigFromFile(model_yml_item->localPath.string());
     auto mc = yaml_handler.GetModelConfig();
     mc.model = unique_model_id;
+
+    uint64_t model_size = 0;
+    for (const auto& item : finishedTask.items) {
+      model_size = model_size + item.bytes.value_or(0);
+    }
+    mc.size = model_size;
     yaml_handler.UpdateModelConfig(mc);
     yaml_handler.WriteYamlFile(model_yml_item->localPath.string());
 
diff --git a/engine/services/model_service.h b/engine/services/model_service.h
index a29f092d4..be450fb0b 100644
--- a/engine/services/model_service.h
+++ b/engine/services/model_service.h
@@ -5,7 +5,6 @@
 #include <string>
 #include "common/engine_servicei.h"
 #include "config/model_config.h"
-#include "database/models.h"
 #include "services/download_service.h"
 #include "services/inference_service.h"