update code

S2I/modules/utils.py

@@ -55,6 +55,31 @@ def downloading(url, outf):
             print("ERROR, something went wrong")
         print(f"Downloaded successfully to {outf}")
 
+def initialize_folder() -> None:
+    """
+    Initialize the folder for storing model weights.
+
+    Raises:
+        OSError: if the folder cannot be created.
+    """
+    home = get_s2i_home()
+    s2i_home_path = home + "/.s2i"
+    weights_path = s2i_home_path + "/weights"
+    print(weights_path)
+    if not os.path.exists(s2i_home_path):
+        os.makedirs(s2i_home_path, exist_ok=True)
+
+    if not os.path.exists(weights_path):
+        os.makedirs(weights_path, exist_ok=True)
+
+def get_s2i_home() -> str:
+    """
+    Get the home directory for storing model weights
+
+    Returns:
+        str: the home directory.
+    """
+    return str(os.getenv("S2I_HOME", default=str(Path.home())))
 
 def download_models():
     urls = {
@@ -62,12 +87,10 @@ def download_models():
         '100k': 'https://huggingface.co/myn0908/sk2ks/resolve/main/model_16001.pkl?download=true',
     }
     # Get the current working directory
-    ckpt_folder = os.path.join(os.getcwd(), 'checkpoints')
-    os.makedirs(ckpt_folder, exist_ok=True)
-
+    home = get_s2i_home()
     model_paths = {}
     for model_name, url in urls.items():
-        outf = os.path.join(ckpt_folder, f"sketch2image_lora_{model_name}.pkl")
+        outf = os.path.join(home, f"sketch2image_lora_{model_name}.pkl")
         downloading(url, outf)
         model_paths[model_name] = outf