set torch threads

app.py

@@ -15,7 +15,7 @@ from interfaces.emotion9 import demo as e9_demo
 from interfaces.cap_media_demo import demo as cap_media_demo
 from interfaces.cap_media2 import demo as cap_media2_demo
 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
+from utils import download_hf_models, hf_cleanup, df_h, set_hf_cache_dir, scan_cache, set_torch_threads
 
 
 css = """
@@ -52,6 +52,7 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
+    set_torch_threads()
     set_hf_cache_dir("/data")
     with hf_cleanup():
         df_h() # debug -> check disk space before launching demo - TO-DO: smarter disk space usage

utils.py

@@ -186,6 +186,12 @@ def set_hf_cache_dir(path:str):
     os.environ['HF_HOME'] = path
     os.environ['HF_DATASETS_CACHE'] = path
     os.environ['TORCH_HOME'] = path
+  
+    
+def set_torch_threads():
+    torch.set_num_threads(1)
+    os.environ["OMP_NUM_THREADS"] = "1"
+    os.environ["MKL_NUM_THREADS"] = "1"
 
 
 def is_disk_full(min_free_space_in_GB=10):