remaining disk space warning

shark/examples/shark_inference/stable_diffusion/main.py

@@ -42,7 +42,7 @@ if args.clear_all:
         shutil.rmtree(os.path.join(home, ".local/shark_tank"))
 
 
-from utils import set_init_device_flags
+from utils import set_init_device_flags, disk_space_check
 
 from opt_params import get_unet, get_vae, get_clip
 from schedulers import (
@@ -104,6 +104,7 @@ if __name__ == "__main__":
         sys.exit("prompts and negative prompts must be of same length")
 
     set_init_device_flags()
+    disk_space_check(Path.cwd())
     clip = get_clip()
     unet = get_unet()
     vae = get_vae()
@@ -260,6 +261,7 @@ if __name__ == "__main__":
         output_path.mkdir(parents=True, exist_ok=True)
     else:
         output_path = Path.cwd()
+    disk_space_check(output_path, lim=5)
     for i in range(batch_size):
         json_store = {
             "prompt": args.prompts[i],

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -242,3 +242,12 @@ def get_available_devices():
     available_devices.extend(cuda_devices)
     available_devices.append("cpu")
     return available_devices
+
+
+def disk_space_check(path, lim=20):
+    from shutil import disk_usage
+
+    du = disk_usage(path)
+    free = du.free / (1024 * 1024 * 1024)
+    if free <= lim:
+        print(f"[WARNING] Only {free:.2f}GB space available in {path}.")

web/models/stable_diffusion/main.py

@@ -5,6 +5,7 @@ import torchvision.transforms as T
 from tqdm.auto import tqdm
 from models.stable_diffusion.cache_objects import model_cache
 from models.stable_diffusion.stable_args import args
+from models.stable_diffusion.utils import disk_space_check
 from random import randint
 import numpy as np
 import time
@@ -72,6 +73,7 @@ def set_ui_params(
 # save output images and the inputs correspoding to it.
 def save_output_img(output_img):
     output_path = args.output_dir if args.output_dir else Path.cwd()
+    disk_space_check(output_path, lim=5)
     generated_imgs_path = Path(output_path, "generated_imgs")
     generated_imgs_path.mkdir(parents=True, exist_ok=True)
     csv_path = Path(generated_imgs_path, "imgs_history.csv")
@@ -140,6 +142,7 @@ def stable_diff_inf(
         width = 768
 
     # get all cached data.
+    disk_space_check(Path.cwd())
     model_cache.set_models(device_key)
     tokenizer = model_cache.tokenizer
     scheduler = model_cache.schedulers[args.scheduler]

web/models/stable_diffusion/utils.py

@@ -232,3 +232,12 @@ def get_available_devices():
     available_devices.extend(cuda_devices)
     #  available_devices.append("cpu")
     return available_devices
+
+
+def disk_space_check(path, lim=20):
+    from shutil import disk_usage
+
+    du = disk_usage(path)
+    free = du.free / (1024 * 1024 * 1024)
+    if free <= lim:
+        print(f"[WARNING] Only {free:.2f}GB space available in {path}.")