[SD] Fix custom model path for WebUI

-- This commit fixes custom model path for WebUI.

Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>

apps/stable_diffusion/web/index.py

@@ -16,7 +16,9 @@
 # Clear all gradio tmp images from the last session
 clear_gradio_tmp_imgs_folder()
 # Create the custom model folder if it doesn't already exist
-get_custom_model_path().mkdir(parents=True, exist_ok=True)
+dir = ["models", "vae", "lora"]
+for root in dir:
+    get_custom_model_path(root).mkdir(parents=True, exist_ok=True)
 
 if args.clear_all:
     clear_all()

apps/stable_diffusion/web/ui/utils.py

@@ -75,14 +75,20 @@ def resource_path(relative_path):
 
 
 def get_custom_model_path(model="models"):
-    prefix = Path(args.ckpt_dir) if args.ckpt_dir else Path(Path.cwd())
+    # If `--ckpt_dir` is provided it'd override the heirarchical folder
+    # structure in WebUI :-
+    #       model
+    #         |___lora
+    #         |___vae
+    if args.ckpt_dir:
+        return Path(args.ckpt_dir)
     match model:
         case "models":
-            return Path(prefix, "models")
+            return Path(Path.cwd(), "models")
         case "vae":
-            return Path(prefix, "models/vae")
+            return Path(Path.cwd(), "models/vae")
         case "lora":
-            return Path(prefix, "models/lora")
+            return Path(Path.cwd(), "models/lora")
         case _:
             return ""