Access to model directory

.gitignore

@@ -3,3 +3,4 @@ venv
 
 __pycache__
 
+model

Dockerfile

@@ -28,6 +28,9 @@ WORKDIR /app
 # Install runtime dependency: libopenblas.so.0 is provided by libopenblas-base.
 RUN apt-get update && apt-get install -y build-essential cmake libopenblas-dev liblapack-dev libopenblas-dev liblapack-dev
 
+# Create the "model" folder for caching the model file
+RUN mkdir -p /app/model
+
 COPY --from=builder /app/venv venv
 
 COPY app.py .

app.py

@@ -7,8 +7,14 @@ import gradio as gr
 
 from huggingface_hub import hf_hub_download
 
-MODEL_PATH = hf_hub_download(repo_id="Robys01/face-aging", filename="best_unet_model.pth")
-print(f"Model downloaded to {MODEL_PATH}")
+MODEL_PATH = "model/best_unet_model.pth"
+
+os.makedirs("model", exist_ok=True)
+
+if not os.path.exists(MODEL_PATH):
+    print("Downloading model from Hugging Face...")
+    path = hf_hub_download(repo_id="Robys01/face-aging", filename="best_unet_model.pth", local_dir="model", cache_dir="model")
+    print(f"Model downloaded to {path}")
 
 model = UNet()  
 model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu"), weights_only=False))