add further caching for torch, hf cache datasets

Dockerfile

@@ -49,7 +49,7 @@ RUN pip install -r requirements.txt
 COPY . .
 
 # make persistent cache for models 
-RUN mkdir -p /app/hf_cache
+RUN mkdir -p /app/hf_cache /app/torch_cache /app/hf_cache/datasets
 
 # Set environment variable to mitigate CUDA memory fragmentation
 ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Image-Morpher/main.py

@@ -13,6 +13,9 @@ import gc
 
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["DIFFUSERS_CACHE"] = "/app/hf_cache"
+os.environ["TORCH_HOME"] = "/app/torch_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
+os.environ["HF_DATASETS_CACHE"] = "/app/hf_cache/datasets"
 
 logs_folder = "logs"
 os.makedirs(logs_folder, exist_ok=True)