cache cleanup

app.py

@@ -4,13 +4,9 @@ from diffusers import StableVideoDiffusionPipeline
 from diffusers.utils import load_image, export_to_video
 import spaces
 import os
+import cleanup  # Import the cleanup script
 
-# Check disk usage
-os.system("df -h")
-
-# Clear Hugging Face cache
-os.system("rm -rf ~/.cache/huggingface")
-os.system("rm -rf ~/.cache/transformers")
+cleanup.free_up_space()  # Call the cleanup function
 
 # Check if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"

cleanup.py

@@ -0,0 +1,23 @@
+import os
+import shutil
+
+def free_up_space():
+    # Check disk usage before cleanup
+    print("Disk usage before cleanup:")
+    os.system("df -h")
+
+    # Clear Hugging Face cache
+    print("Clearing Hugging Face cache...")
+    shutil.rmtree(os.path.expanduser("~/.cache/huggingface"), ignore_errors=True)
+    shutil.rmtree(os.path.expanduser("~/.cache/transformers"), ignore_errors=True)
+
+    # Clear PyTorch cache
+    print("Clearing PyTorch cache...")
+    shutil.rmtree(os.path.expanduser("~/.cache/torch"), ignore_errors=True)
+
+    # Check disk usage after cleanup
+    print("Disk usage after cleanup:")
+    os.system("df -h")
+
+if __name__ == "__main__":
+    free_up_space()