Model: Fix load packets

The model_type internal reference was changed to an enum for a more extendable loading process. Return the current model type when loading a new model. Signed-off-by: kingbri <[email protected]>
theroyallab · Nov 21, 2024 · c49047e · c49047e
1 parent 0ab393f
commit c49047e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/common/model.py b/common/model.py
@@ -88,14 +88,15 @@ async def load_model_gen(model_path: pathlib.Path, **kwargs):
     try:
         index = 0
         async for module, modules in load_status:
+            current_model_type = model_type[index].value
             if module == 0:
                 loading_task = progress.add_task(
-                    f"[cyan]Loading {model_type[index].value} modules", total=modules
+                    f"[cyan]Loading {current_model_type} modules", total=modules
                 )
             else:
                 progress.advance(loading_task)
 
-            yield module, modules, model_type
+            yield module, modules, current_model_type
 
             if module == modules:
                 # Switch to model progress if the draft model is loaded