Main: Don't load if model_name/loras is blank

Previously, if model_name was commented out, a load would not occur. Add the case if model_name or loras is blank which returns None when parsing the YAML. Signed-off-by: kingbri <[email protected]>
theroyallab · Jan 2, 2024 · 451042a · 451042a
1 parent 6b04463
commit 451042a
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/main.py b/main.py
@@ -525,9 +525,10 @@ def entrypoint(args: Optional[dict] = None):
     # If an initial model name is specified, create a container
     # and load the model
     model_config = get_model_config()
-    if "model_name" in model_config:
+    model_name = model_config.get("model_name")
+    if model_name:
         model_path = pathlib.Path(unwrap(model_config.get("model_dir"), "models"))
-        model_path = model_path / model_config.get("model_name")
+        model_path = model_path / model_name
 
         MODEL_CONTAINER = ModelContainer(model_path.resolve(), False, **model_config)
         load_status = MODEL_CONTAINER.load_gen(load_progress)
@@ -542,7 +543,7 @@ def entrypoint(args: Optional[dict] = None):
 
         # Load loras after loading the model
         lora_config = get_lora_config()
-        if "loras" in lora_config:
+        if lora_config.get("loras"):
             lora_dir = pathlib.Path(unwrap(lora_config.get("lora_dir"), "loras"))
             MODEL_CONTAINER.load_loras(lora_dir.resolve(), **lora_config)