diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py index 78b9295e8..8d45327ef 100644 --- a/modules/civitai/download_civitai.py +++ b/modules/civitai/download_civitai.py @@ -1,4 +1,5 @@ import os +import re import uuid import hashlib import threading @@ -12,6 +13,10 @@ from modules import shared, paths from modules.logger import console +STALE_PARTIAL_DAYS = 7 +TEMP_FILE_RE = re.compile(r'^[0-9a-f]{8}\.tmp$') + + @dataclass class DownloadItem: id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) @@ -57,6 +62,27 @@ class DownloadManager: self._lock = threading.Lock() self._max_workers = max_workers self._worker_count = 0 + threading.Thread(target=self._sweep_stale_partials, daemon=True).start() + + def _sweep_stale_partials(self): + # Partials double as resume state, but mtime stops moving the moment a + # download stops, so an untouched week-old .tmp is abandoned. + try: + from modules.civitai.filemanage_civitai import iter_type_roots + cutoff = time.time() - STALE_PARTIAL_DAYS * 86400 + for root in iter_type_roots(): + for path in root.rglob('*.tmp'): + if not TEMP_FILE_RE.match(path.name): + continue + try: + stat = path.stat() + if stat.st_mtime < cutoff: + path.unlink() + log.info(f'CivitAI stale partial removed: file="{path}" size={stat.st_size/1024/1024:.0f}MB age>{STALE_PARTIAL_DAYS}d') + except OSError as e: + log.warning(f'CivitAI stale partial sweep: file="{path}" {e}') + except Exception as e: + log.warning(f'CivitAI stale partial sweep error: {e}') def enqueue(self, url: str, folder: str, filename: str, model_type: str = "", expected_hash: str = "", token: str | None = None, diff --git a/modules/civitai/filemanage_civitai.py b/modules/civitai/filemanage_civitai.py index a914d0611..c889a107a 100644 --- a/modules/civitai/filemanage_civitai.py +++ b/modules/civitai/filemanage_civitai.py @@ -60,6 +60,25 @@ def get_type_folder(model_type: str, base_model: str = '') -> Path: return Path(paths.models_path) / fallback_dir +def iter_type_roots() -> set[Path]: + """Every root folder downloads can resolve into, for maintenance sweeps.""" + from modules import shared, paths + roots = set() + custom_json = getattr(shared.opts, 'civitai_save_type_folders', '') or '' + if custom_json.strip(): + try: + import json + for folder in json.loads(custom_json).values(): + p = Path(folder) + roots.add(p if p.is_absolute() else Path(paths.models_path) / folder) + except Exception: + pass + for opt_attr, fallback_dir in set(TYPE_MAP.values()) | {('unet_dir', 'UNET')}: + configured = (getattr(shared.opts, opt_attr, '') or '') if opt_attr else '' + roots.add(Path(configured) if configured else Path(paths.models_path) / fallback_dir) + return {r for r in roots if r.is_dir()} + + def resolve_save_path(model_type: str, model_name: str = "", base_model: str = "", nsfw: bool = False, creator: str = "", model_id: int = 0, version_id: int = 0, version_name: str = "") -> Path: