save to proper storage not ephemeral

gradio_app.py

@@ -3,38 +3,42 @@ import torch
 import os
 import time
 import datetime
+from pathlib import Path
 from moviepy.editor import VideoFileClip
 import gradio as gr
 
 # Download Weights
 from huggingface_hub import snapshot_download
 
-# List of subdirectories to create inside "weights"
+# Decide where to cache weights. Default to persistent disk (/data) when available, fall back to workspace.
+default_weights_root = Path("/data/weights") if Path("/data").exists() else Path("./weights")
+WEIGHTS_ROOT = Path(os.environ.get("DIFFUERASER_WEIGHTS_ROOT", default_weights_root))
+
+# List of subdirectories to create inside the weights root
 subfolders = [
     "diffuEraser",
     "stable-diffusion-v1-5",
     "PCM_Weights",
     "propainter",
-    "sd-vae-ft-mse"
+    "sd-vae-ft-mse",
 ]
-# Create directories
 for subfolder in subfolders:
-    os.makedirs(os.path.join("weights", subfolder), exist_ok=True)
+    (WEIGHTS_ROOT / subfolder).mkdir(parents=True, exist_ok=True)
 
-snapshot_download(repo_id="lixiaowen/diffuEraser", local_dir="./weights/diffuEraser")
-snapshot_download(repo_id="stable-diffusion-v1-5/stable-diffusion-v1-5", local_dir="./weights/stable-diffusion-v1-5")
-snapshot_download(repo_id="wangfuyun/PCM_Weights", local_dir="./weights/PCM_Weights")
-snapshot_download(repo_id="camenduru/ProPainter", local_dir="./weights/propainter")
-snapshot_download(repo_id="stabilityai/sd-vae-ft-mse", local_dir="./weights/sd-vae-ft-mse")
+snapshot_download(repo_id="lixiaowen/diffuEraser", local_dir=str(WEIGHTS_ROOT / "diffuEraser"))
+snapshot_download(repo_id="stable-diffusion-v1-5/stable-diffusion-v1-5", local_dir=str(WEIGHTS_ROOT / "stable-diffusion-v1-5"))
+snapshot_download(repo_id="wangfuyun/PCM_Weights", local_dir=str(WEIGHTS_ROOT / "PCM_Weights"))
+snapshot_download(repo_id="camenduru/ProPainter", local_dir=str(WEIGHTS_ROOT / "propainter"))
+snapshot_download(repo_id="stabilityai/sd-vae-ft-mse", local_dir=str(WEIGHTS_ROOT / "sd-vae-ft-mse"))
 
 # Import model classes
 from diffueraser.diffueraser import DiffuEraser
 from propainter.inference import Propainter, get_device
 
-base_model_path = "weights/stable-diffusion-v1-5"
-vae_path = "weights/sd-vae-ft-mse"
-diffueraser_path = "weights/diffuEraser"
-propainter_model_dir = "weights/propainter"
+base_model_path = str(WEIGHTS_ROOT / "stable-diffusion-v1-5")
+vae_path = str(WEIGHTS_ROOT / "sd-vae-ft-mse")
+diffueraser_path = str(WEIGHTS_ROOT / "diffuEraser")
+propainter_model_dir = str(WEIGHTS_ROOT / "propainter")
 
 # Model setup
 device = get_device()