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 <bdashore3@proton.me>
This commit is contained in:
kingbri 2024-11-21 18:06:47 -05:00
parent 0ab393f09c
commit c49047eea1

View file

@ -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