mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
feat(civitai): sweep abandoned download partials at startup
Interrupted downloads leave url-hash .tmp resume files that nothing cleans up; the queue that knows about them is in-memory only. - iter_type_roots lists every folder downloads resolve into - partials older than 7 days by mtime are deleted when the download manager starts; active partials always have a fresh mtime
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user