hf cleanup as context

app.py

@@ -18,7 +18,6 @@ from interfaces.cap_minor_media import demo as cap_minor_media_demo
 from utils import download_hf_models, hf_cleanup, df_h, set_hf_cache_dir, scan_cache
 
 
-
 css = """
 /* Make only the active tab bold */
 .svelte-1uw5tnk[aria-selected="true"] {
@@ -54,12 +53,11 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     set_hf_cache_dir("/data")
-    hf_cleanup()
-    df_h() # debug -> check disk space before launching demo - TO-DO: smarter disk space usage
-    scan_cache()
-    download_spacy_models() 
-    download_hf_models() # does this affect the build?
-    hf_cleanup()
+    with hf_cleanup():
+        df_h() # debug -> check disk space before launching demo - TO-DO: smarter disk space usage
+        scan_cache()
+        download_spacy_models() 
+        download_hf_models() # does this affect the build?
     demo.launch()
 
 # TODO: add all languages & domains

utils.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import glob
 import subprocess
+from contextlib import contextmanager
 
 import torch
 
@@ -116,12 +117,21 @@ def df_h():
     print(du_result.stdout)
     
     
-def hf_cleanup():
+def delete_http_folders():
     http_folders = glob.glob("/data/http*")
     for folder in http_folders:
         if os.path.isdir(folder):
             print(f"Deleting: {folder}")
             shutil.rmtree(folder)
+       
+            
+@contextmanager
+def hf_cleanup():
+    delete_http_folders()
+    try:
+        yield
+    finally:
+        delete_http_folders()
     
     
 def scan_cache():