Add stack Update, ComfyUI model manager, and slim ComfyUI stack.
Server UI gains Update on stack cards, ComfyUI Models tab with workflow scan and downloads, and centralized comfyui_config. Model catalog and download scripts move from stacks/comfyui to server-ui so ComfyUI stays a minimal Docker wrapper for easier image updates.
This commit is contained in:
tomasz-syn-grzegorza
@@ -0,0 +1,858 @@
|
||||
"""ComfyUI model catalog and download jobs for Server UI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import yaml
|
||||
|
||||
from comfyui_config import (
|
||||
COMFYUI_CONTAINER,
|
||||
catalog_path as comfyui_catalog_path,
|
||||
host_models_root,
|
||||
host_to_container_model_path,
|
||||
host_workflows_dir,
|
||||
scripts_dir as comfyui_scripts_dir,
|
||||
stack_env_path as comfyui_stack_env_path,
|
||||
)
|
||||
from comfyui_manager_catalog import load_manager_index
|
||||
from workflow_scanner import find_model_file, scan_workflows
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
WORKFLOW_MODEL_PREFIX = "wf:"
|
||||
|
||||
CATEGORY_DIRS: dict[str, str] = {
|
||||
"checkpoint": "checkpoints",
|
||||
"lora": "loras",
|
||||
"vae": "vae",
|
||||
"controlnet": "controlnet",
|
||||
"upscale": "upscale_models",
|
||||
"latent_upscale": "latent_upscale_models",
|
||||
"text_encoder": "text_encoders",
|
||||
"diffusion_models": "diffusion_models",
|
||||
"clip": "clip",
|
||||
"unet": "unet",
|
||||
"unknown": "checkpoints",
|
||||
}
|
||||
|
||||
LOG_BUFFER_MAX = 256 * 1024
|
||||
JOB_TTL_SECONDS = 3600
|
||||
PROGRESS_RE = re.compile(r"(\d{1,3})%")
|
||||
|
||||
_job_manager: DownloadJobManager | None = None
|
||||
|
||||
|
||||
class ComfyUIModelError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ModelNotFoundError(ComfyUIModelError):
|
||||
pass
|
||||
|
||||
|
||||
class JobNotFoundError(ComfyUIModelError):
|
||||
pass
|
||||
|
||||
|
||||
class JobConflictError(ComfyUIModelError):
|
||||
pass
|
||||
|
||||
|
||||
class DeleteModelError(ComfyUIModelError):
|
||||
pass
|
||||
|
||||
|
||||
def models_root(data_root: Path) -> Path:
|
||||
return host_models_root(data_root)
|
||||
|
||||
|
||||
def workflows_dir(data_root: Path) -> Path:
|
||||
return host_workflows_dir(data_root)
|
||||
|
||||
|
||||
def workflow_model_id(filename: str) -> str:
|
||||
return f"{WORKFLOW_MODEL_PREFIX}{filename}"
|
||||
|
||||
|
||||
def is_workflow_model_id(model_id: str) -> bool:
|
||||
return model_id.startswith(WORKFLOW_MODEL_PREFIX)
|
||||
|
||||
|
||||
def resolve_relative_model_path(relative_path: str, data_root: Path) -> Path:
|
||||
rel = relative_path.strip().lstrip("/")
|
||||
if not rel or ".." in Path(rel).parts:
|
||||
raise DeleteModelError(f"Invalid relative model path: {relative_path}")
|
||||
dest = models_root(data_root) / rel
|
||||
if not _is_under_models_root(dest, data_root):
|
||||
raise DeleteModelError(f"Path outside models root: {dest}")
|
||||
return dest
|
||||
|
||||
|
||||
def _validate_download_url(url: str) -> None:
|
||||
parsed = urlparse(url.strip())
|
||||
if parsed.scheme not in ("http", "https") or not parsed.netloc:
|
||||
raise ComfyUIModelError("URL must start with http:// or https://")
|
||||
|
||||
|
||||
def _catalog_by_filename(catalog_path: Path) -> dict[str, dict[str, Any]]:
|
||||
by_filename: dict[str, dict[str, Any]] = {}
|
||||
for entry in load_catalog(catalog_path):
|
||||
filename = entry.get("filename")
|
||||
model_id = entry.get("id")
|
||||
if filename and model_id:
|
||||
by_filename[str(filename)] = entry
|
||||
return by_filename
|
||||
|
||||
|
||||
def build_workflow_required(
|
||||
catalog_path: Path,
|
||||
data_root: Path,
|
||||
active_by_model: dict[str, str],
|
||||
jobs: dict[str, DownloadJob] | None = None,
|
||||
) -> tuple[list[dict[str, Any]], int, str]:
|
||||
"""Build workflow-required model list, excluding catalog entries (dedup)."""
|
||||
jobs = jobs or {}
|
||||
wf_dir = workflows_dir(data_root)
|
||||
refs, workflows_scanned = scan_workflows(wf_dir)
|
||||
catalog_filenames = {
|
||||
str(e.get("filename", "")) for e in load_catalog(catalog_path) if e.get("filename")
|
||||
}
|
||||
manager_index = load_manager_index(data_root)
|
||||
root = models_root(data_root)
|
||||
items: list[dict[str, Any]] = []
|
||||
|
||||
for ref in refs:
|
||||
if ref.filename in catalog_filenames:
|
||||
continue
|
||||
|
||||
model_id = workflow_model_id(ref.filename)
|
||||
found = find_model_file(root, ref.filename)
|
||||
|
||||
source = "none"
|
||||
download_url: str | None = None
|
||||
save_path: str | None = None
|
||||
catalog_id: str | None = None
|
||||
size_gb: Any = None
|
||||
note = ""
|
||||
|
||||
catalog_entry = _catalog_by_filename(catalog_path).get(ref.filename)
|
||||
if catalog_entry:
|
||||
source = "catalog"
|
||||
catalog_id = str(catalog_entry.get("id"))
|
||||
download_url = catalog_entry.get("download_url")
|
||||
size_gb = catalog_entry.get("size_gb")
|
||||
note = str(catalog_entry.get("note") or "")
|
||||
dest = resolve_dest_path(catalog_entry, data_root)
|
||||
else:
|
||||
mgr = manager_index.get(ref.filename)
|
||||
if mgr and mgr.get("url"):
|
||||
source = "manager"
|
||||
download_url = str(mgr["url"])
|
||||
save_path = str(mgr.get("save_path") or category_dir(ref.category))
|
||||
size_str = mgr.get("size")
|
||||
if isinstance(size_str, str) and size_str.upper().endswith("GB"):
|
||||
try:
|
||||
size_gb = float(size_str.upper().replace("GB", "").strip())
|
||||
except ValueError:
|
||||
size_gb = None
|
||||
note = str(mgr.get("description") or mgr.get("name") or "")
|
||||
dest = root / save_path / ref.filename
|
||||
else:
|
||||
dest = root / category_dir(ref.category) / ref.filename
|
||||
|
||||
if found is not None:
|
||||
dest = found
|
||||
|
||||
try:
|
||||
relative_path = dest.relative_to(root).as_posix()
|
||||
except ValueError:
|
||||
relative_path = None
|
||||
active_job: DownloadJob | None = None
|
||||
job_id = active_by_model.get(model_id)
|
||||
if job_id and job_id in jobs:
|
||||
active_job = jobs[job_id]
|
||||
|
||||
state = _model_file_state(dest, size_gb, active_job)
|
||||
items.append(
|
||||
{
|
||||
"id": model_id,
|
||||
"name": ref.filename,
|
||||
"filename": ref.filename,
|
||||
"category": ref.category,
|
||||
"workflows": ref.workflows,
|
||||
"source": source,
|
||||
"downloadable": source in ("catalog", "manager"),
|
||||
"download_url": download_url,
|
||||
"save_path": save_path,
|
||||
"catalog_id": catalog_id,
|
||||
"relative_path": relative_path,
|
||||
"note": note,
|
||||
"size_gb": size_gb,
|
||||
**state,
|
||||
}
|
||||
)
|
||||
|
||||
return items, workflows_scanned, str(wf_dir)
|
||||
|
||||
|
||||
def _expected_bytes(size_gb: Any) -> int | None:
|
||||
if size_gb is None:
|
||||
return None
|
||||
try:
|
||||
return int(float(size_gb) * 1_000_000_000)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _is_under_models_root(path: Path, data_root: Path) -> bool:
|
||||
root = models_root(data_root).resolve()
|
||||
try:
|
||||
path.resolve().relative_to(root)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _file_size(path: Path) -> int:
|
||||
if not path.is_file():
|
||||
return 0
|
||||
return path.stat().st_size
|
||||
|
||||
|
||||
def category_dir(category: str) -> str:
|
||||
return CATEGORY_DIRS.get(category, category)
|
||||
|
||||
|
||||
def resolve_dest_path(entry: dict[str, Any], data_root: Path) -> Path:
|
||||
local_path = entry.get("local_path")
|
||||
if local_path:
|
||||
return Path(local_path)
|
||||
category = entry.get("category", "")
|
||||
filename = entry.get("filename", "")
|
||||
subdir = category_dir(str(category))
|
||||
return data_root / "apps" / "comfyui" / "models" / subdir / filename
|
||||
|
||||
|
||||
def load_catalog(catalog_path: Path) -> list[dict[str, Any]]:
|
||||
if not catalog_path.is_file():
|
||||
raise ComfyUIModelError(f"Catalog not found: {catalog_path}")
|
||||
with catalog_path.open(encoding="utf-8") as fh:
|
||||
data = yaml.safe_load(fh) or {}
|
||||
models = data.get("models") or []
|
||||
if not isinstance(models, list):
|
||||
raise ComfyUIModelError("Invalid catalog: models must be a list")
|
||||
return models
|
||||
|
||||
|
||||
def _model_file_state(
|
||||
dest: Path,
|
||||
size_gb: Any,
|
||||
active_job: DownloadJob | None,
|
||||
) -> dict[str, Any]:
|
||||
size_bytes = _file_size(dest)
|
||||
expected = _expected_bytes(size_gb)
|
||||
download_info: dict[str, Any] | None = None
|
||||
if active_job and active_job.status == "running":
|
||||
download_info = {
|
||||
"job_id": active_job.id,
|
||||
"status": active_job.status,
|
||||
"progress": active_job.progress,
|
||||
}
|
||||
|
||||
installed = False
|
||||
partial = False
|
||||
if size_bytes > 0:
|
||||
if active_job and active_job.status == "running":
|
||||
partial = True
|
||||
elif expected is None or size_bytes >= int(expected * 0.95):
|
||||
installed = True
|
||||
else:
|
||||
partial = True
|
||||
|
||||
return {
|
||||
"path": str(dest),
|
||||
"size_bytes": size_bytes,
|
||||
"expected_bytes": expected,
|
||||
"installed": installed,
|
||||
"partial": partial,
|
||||
"download": download_info,
|
||||
}
|
||||
|
||||
|
||||
def _unlink_model_file(dest: Path, data_root: Path) -> bool:
|
||||
if not dest.is_file():
|
||||
return False
|
||||
if not _is_under_models_root(dest, data_root):
|
||||
raise DeleteModelError(f"Refusing to delete path outside models root: {dest}")
|
||||
|
||||
try:
|
||||
dest.unlink()
|
||||
if not dest.is_file():
|
||||
return True
|
||||
except OSError as exc:
|
||||
log.warning("Host unlink failed for %s: %s", dest, exc)
|
||||
|
||||
container_path = host_to_container_model_path(dest, data_root)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "exec", COMFYUI_CONTAINER, "rm", "-f", container_path],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
log.warning("docker exec rm failed: %s", result.stderr.strip())
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
log.warning("docker exec rm error for %s: %s", container_path, exc)
|
||||
|
||||
return not dest.is_file()
|
||||
|
||||
|
||||
def list_models(catalog_path: Path, data_root: Path) -> list[dict[str, Any]]:
|
||||
"""Basic catalog list without active download jobs (legacy helper)."""
|
||||
return _build_model_list(catalog_path, data_root, active_by_model={})
|
||||
|
||||
|
||||
def _build_model_list(
|
||||
catalog_path: Path,
|
||||
data_root: Path,
|
||||
active_by_model: dict[str, str],
|
||||
jobs: dict[str, DownloadJob] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
jobs = jobs or {}
|
||||
items: list[dict[str, Any]] = []
|
||||
for entry in load_catalog(catalog_path):
|
||||
model_id = entry.get("id")
|
||||
if not model_id:
|
||||
continue
|
||||
dest = resolve_dest_path(entry, data_root)
|
||||
active_job: DownloadJob | None = None
|
||||
job_id = active_by_model.get(str(model_id))
|
||||
if job_id and job_id in jobs:
|
||||
active_job = jobs[job_id]
|
||||
state = _model_file_state(dest, entry.get("size_gb"), active_job)
|
||||
items.append(
|
||||
{
|
||||
"id": model_id,
|
||||
"name": entry.get("name", model_id),
|
||||
"category": entry.get("category", ""),
|
||||
"source": entry.get("source", ""),
|
||||
"size_gb": entry.get("size_gb"),
|
||||
"note": entry.get("note", ""),
|
||||
"filename": entry.get("filename", dest.name),
|
||||
**state,
|
||||
}
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
def list_installed(data_root: Path) -> list[dict[str, Any]]:
|
||||
models_root = data_root / "apps" / "comfyui" / "models"
|
||||
if not models_root.is_dir():
|
||||
return []
|
||||
installed: list[dict[str, Any]] = []
|
||||
for path in sorted(models_root.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
if path.name.startswith("."):
|
||||
continue
|
||||
rel = path.relative_to(models_root)
|
||||
installed.append(
|
||||
{
|
||||
"path": str(path),
|
||||
"relative_path": str(rel),
|
||||
"category_dir": rel.parts[0] if rel.parts else "",
|
||||
"filename": path.name,
|
||||
"size_bytes": path.stat().st_size,
|
||||
}
|
||||
)
|
||||
return installed
|
||||
|
||||
|
||||
def disk_usage(data_root: Path) -> dict[str, Any]:
|
||||
usage = shutil.disk_usage(data_root)
|
||||
return {
|
||||
"path": str(data_root),
|
||||
"total_bytes": usage.total,
|
||||
"used_bytes": usage.used,
|
||||
"free_bytes": usage.free,
|
||||
}
|
||||
|
||||
|
||||
def _parse_progress(line: str, current: int) -> int:
|
||||
match = PROGRESS_RE.search(line)
|
||||
if match:
|
||||
return max(current, min(100, int(match.group(1))))
|
||||
lower = line.lower()
|
||||
if "already exists" in lower:
|
||||
return 100
|
||||
if line.strip() == "Done.":
|
||||
return 100
|
||||
return current
|
||||
|
||||
|
||||
@dataclass
|
||||
class DownloadJob:
|
||||
id: str
|
||||
model_id: str
|
||||
status: str = "running"
|
||||
progress: int = 0
|
||||
log: str = ""
|
||||
error: str = ""
|
||||
started_at: float = field(default_factory=time.time)
|
||||
finished_at: float | None = None
|
||||
file_dest: Path | None = None
|
||||
_log_buffer: list[str] = field(default_factory=list)
|
||||
_process: subprocess.Popen[str] | None = field(default=None, repr=False)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
|
||||
|
||||
def append_log(self, chunk: str) -> None:
|
||||
with self._lock:
|
||||
self._log_buffer.append(chunk)
|
||||
combined = "".join(self._log_buffer)
|
||||
if len(combined) > LOG_BUFFER_MAX:
|
||||
combined = combined[-LOG_BUFFER_MAX:]
|
||||
self._log_buffer = [combined]
|
||||
self.log = combined
|
||||
for line in chunk.splitlines():
|
||||
self.progress = _parse_progress(line, self.progress)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"id": self.id,
|
||||
"model_id": self.model_id,
|
||||
"status": self.status,
|
||||
"progress": self.progress,
|
||||
"log": self.log,
|
||||
"error": self.error,
|
||||
"started_at": self.started_at,
|
||||
"finished_at": self.finished_at,
|
||||
}
|
||||
|
||||
|
||||
class DownloadJobManager:
|
||||
def __init__(
|
||||
self,
|
||||
catalog_path: Path,
|
||||
data_root: Path,
|
||||
download_script: Path,
|
||||
download_file_script: Path,
|
||||
stack_env_path: Path | None = None,
|
||||
job_ttl: int = JOB_TTL_SECONDS,
|
||||
) -> None:
|
||||
self.catalog_path = catalog_path
|
||||
self.data_root = data_root
|
||||
self.download_script = download_script
|
||||
self.download_file_script = download_file_script
|
||||
self.stack_env_path = stack_env_path
|
||||
self.job_ttl = job_ttl
|
||||
self._jobs: dict[str, DownloadJob] = {}
|
||||
self._active_by_model: dict[str, str] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _catalog_by_id(self) -> dict[str, dict[str, Any]]:
|
||||
return {str(m["id"]): m for m in load_catalog(self.catalog_path) if m.get("id")}
|
||||
|
||||
def _cleanup_old_jobs(self) -> None:
|
||||
cutoff = time.time() - self.job_ttl
|
||||
with self._lock:
|
||||
expired = [
|
||||
jid
|
||||
for jid, job in self._jobs.items()
|
||||
if job.finished_at is not None and job.finished_at < cutoff
|
||||
]
|
||||
for jid in expired:
|
||||
job = self._jobs.pop(jid, None)
|
||||
if job and job.model_id in self._active_by_model:
|
||||
if self._active_by_model.get(job.model_id) == jid:
|
||||
del self._active_by_model[job.model_id]
|
||||
|
||||
def get_job(self, job_id: str) -> DownloadJob:
|
||||
self._cleanup_old_jobs()
|
||||
job = self._jobs.get(job_id)
|
||||
if not job:
|
||||
raise JobNotFoundError(f"Unknown job: {job_id}")
|
||||
return job
|
||||
|
||||
def list_models_with_state(self) -> list[dict[str, Any]]:
|
||||
self._cleanup_old_jobs()
|
||||
with self._lock:
|
||||
active = dict(self._active_by_model)
|
||||
jobs = dict(self._jobs)
|
||||
return _build_model_list(self.catalog_path, self.data_root, active, jobs)
|
||||
|
||||
def list_models_response(self) -> dict[str, Any]:
|
||||
self._cleanup_old_jobs()
|
||||
with self._lock:
|
||||
active = dict(self._active_by_model)
|
||||
jobs = dict(self._jobs)
|
||||
models = _build_model_list(self.catalog_path, self.data_root, active, jobs)
|
||||
workflow_required, workflows_scanned, workflows_dir_path = build_workflow_required(
|
||||
self.catalog_path, self.data_root, active, jobs
|
||||
)
|
||||
return {
|
||||
"models": models,
|
||||
"workflow_required": workflow_required,
|
||||
"workflows_scanned": workflows_scanned,
|
||||
"workflows_dir": workflows_dir_path,
|
||||
}
|
||||
|
||||
def _dest_for_model(self, model_id: str) -> Path:
|
||||
catalog = self._catalog_by_id()
|
||||
if model_id not in catalog:
|
||||
raise ModelNotFoundError(f"Unknown model id: {model_id}")
|
||||
dest = resolve_dest_path(catalog[model_id], self.data_root)
|
||||
if not _is_under_models_root(dest, self.data_root):
|
||||
raise DeleteModelError(f"Invalid model path: {dest}")
|
||||
return dest
|
||||
|
||||
def _remove_model_file(self, model_id: str) -> bool:
|
||||
dest = self._dest_for_model(model_id)
|
||||
return _unlink_model_file(dest, self.data_root)
|
||||
|
||||
def delete_model(self, model_id: str) -> dict[str, Any]:
|
||||
self._cleanup_old_jobs()
|
||||
with self._lock:
|
||||
job_id = self._active_by_model.get(model_id)
|
||||
if job_id:
|
||||
self.cancel_job(job_id, cleanup=True)
|
||||
if is_workflow_model_id(model_id):
|
||||
raise ModelNotFoundError(
|
||||
f"Workflow model {model_id} must be deleted via relative_path"
|
||||
)
|
||||
removed = self._remove_model_file(model_id)
|
||||
dest = self._dest_for_model(model_id)
|
||||
return {
|
||||
"model_id": model_id,
|
||||
"removed": removed,
|
||||
"path": str(dest),
|
||||
}
|
||||
|
||||
def delete_model_file(self, relative_path: str) -> dict[str, Any]:
|
||||
self._cleanup_old_jobs()
|
||||
dest = resolve_relative_model_path(relative_path, self.data_root)
|
||||
model_id = workflow_model_id(dest.name)
|
||||
with self._lock:
|
||||
job_id = self._active_by_model.get(model_id)
|
||||
if job_id:
|
||||
self.cancel_job(job_id, cleanup=True)
|
||||
removed = _unlink_model_file(dest, self.data_root)
|
||||
return {
|
||||
"relative_path": relative_path,
|
||||
"removed": removed,
|
||||
"path": str(dest),
|
||||
}
|
||||
|
||||
def start_download(self, model_id: str) -> DownloadJob:
|
||||
self._cleanup_old_jobs()
|
||||
catalog = self._catalog_by_id()
|
||||
if model_id not in catalog:
|
||||
raise ModelNotFoundError(f"Unknown model id: {model_id}")
|
||||
|
||||
entry = catalog[model_id]
|
||||
dest = resolve_dest_path(entry, self.data_root)
|
||||
if dest.is_file() and _model_file_state(dest, entry.get("size_gb"), None)["installed"]:
|
||||
job = DownloadJob(
|
||||
id=str(uuid.uuid4()),
|
||||
model_id=model_id,
|
||||
status="completed",
|
||||
progress=100,
|
||||
log=f"Already exists: {dest}\n",
|
||||
finished_at=time.time(),
|
||||
)
|
||||
with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
return job
|
||||
|
||||
with self._lock:
|
||||
if model_id in self._active_by_model:
|
||||
active_id = self._active_by_model[model_id]
|
||||
active = self._jobs.get(active_id)
|
||||
if active and active.status == "running":
|
||||
raise JobConflictError(
|
||||
f"Download already running for {model_id} (job {active_id})"
|
||||
)
|
||||
|
||||
if not self.download_script.is_file():
|
||||
raise ComfyUIModelError(f"Download script not found: {self.download_script}")
|
||||
|
||||
job = DownloadJob(id=str(uuid.uuid4()), model_id=model_id)
|
||||
with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
self._active_by_model[model_id] = job.id
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._run_download,
|
||||
args=(job,),
|
||||
name=f"comfyui-download-{job.id[:8]}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return job
|
||||
|
||||
def start_file_download(
|
||||
self,
|
||||
url: str,
|
||||
relative_path: str,
|
||||
*,
|
||||
model_id: str | None = None,
|
||||
size_gb: Any = None,
|
||||
) -> DownloadJob:
|
||||
self._cleanup_old_jobs()
|
||||
_validate_download_url(url)
|
||||
dest = resolve_relative_model_path(relative_path, self.data_root)
|
||||
wf_id = model_id or workflow_model_id(dest.name)
|
||||
|
||||
if dest.is_file() and _model_file_state(dest, size_gb, None)["installed"]:
|
||||
job = DownloadJob(
|
||||
id=str(uuid.uuid4()),
|
||||
model_id=wf_id,
|
||||
status="completed",
|
||||
progress=100,
|
||||
log=f"Already exists: {dest}\n",
|
||||
finished_at=time.time(),
|
||||
file_dest=dest,
|
||||
)
|
||||
with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
return job
|
||||
|
||||
with self._lock:
|
||||
if wf_id in self._active_by_model:
|
||||
active_id = self._active_by_model[wf_id]
|
||||
active = self._jobs.get(active_id)
|
||||
if active and active.status == "running":
|
||||
raise JobConflictError(
|
||||
f"Download already running for {wf_id} (job {active_id})"
|
||||
)
|
||||
|
||||
if not self.download_file_script.is_file():
|
||||
raise ComfyUIModelError(
|
||||
f"Download file script not found: {self.download_file_script}"
|
||||
)
|
||||
|
||||
job = DownloadJob(
|
||||
id=str(uuid.uuid4()),
|
||||
model_id=wf_id,
|
||||
file_dest=dest,
|
||||
)
|
||||
with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
self._active_by_model[wf_id] = job.id
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._run_file_download,
|
||||
args=(job, url, dest),
|
||||
name=f"comfyui-file-download-{job.id[:8]}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return job
|
||||
|
||||
def start_workflow_download(self, model_id: str) -> DownloadJob:
|
||||
"""Download a workflow-required model using catalog or manager metadata."""
|
||||
if not is_workflow_model_id(model_id):
|
||||
raise ModelNotFoundError(f"Not a workflow model id: {model_id}")
|
||||
|
||||
with self._lock:
|
||||
active = dict(self._active_by_model)
|
||||
jobs = dict(self._jobs)
|
||||
items, _, _ = build_workflow_required(
|
||||
self.catalog_path, self.data_root, active, jobs
|
||||
)
|
||||
entry = next((m for m in items if m["id"] == model_id), None)
|
||||
if not entry:
|
||||
raise ModelNotFoundError(f"Unknown workflow model: {model_id}")
|
||||
if not entry.get("downloadable"):
|
||||
raise ComfyUIModelError(
|
||||
f"No download URL for {entry.get('filename')} — add to catalog or ComfyUI-Manager"
|
||||
)
|
||||
|
||||
if entry.get("catalog_id"):
|
||||
return self.start_download(str(entry["catalog_id"]))
|
||||
|
||||
relative_path = entry.get("relative_path")
|
||||
download_url = entry.get("download_url")
|
||||
if not relative_path or not download_url:
|
||||
raise ComfyUIModelError(f"Missing download metadata for {model_id}")
|
||||
|
||||
return self.start_file_download(
|
||||
str(download_url),
|
||||
str(relative_path),
|
||||
model_id=model_id,
|
||||
size_gb=entry.get("size_gb"),
|
||||
)
|
||||
|
||||
def cancel_job(self, job_id: str, *, cleanup: bool = True) -> DownloadJob:
|
||||
job = self.get_job(job_id)
|
||||
if job.status == "running":
|
||||
proc = job._process
|
||||
if proc and proc.poll() is None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
job.status = "cancelled"
|
||||
job.finished_at = time.time()
|
||||
job.append_log("\nCancelled by user.\n")
|
||||
with self._lock:
|
||||
if self._active_by_model.get(job.model_id) == job.id:
|
||||
del self._active_by_model[job.model_id]
|
||||
if cleanup:
|
||||
try:
|
||||
if job.file_dest:
|
||||
_unlink_model_file(job.file_dest, self.data_root)
|
||||
job.append_log("Partial file removed.\n")
|
||||
else:
|
||||
self._remove_model_file(job.model_id)
|
||||
job.append_log("Partial file removed.\n")
|
||||
except ComfyUIModelError as exc:
|
||||
job.append_log(f"Cleanup failed: {exc}\n")
|
||||
return job
|
||||
|
||||
def _run_file_download(self, job: DownloadJob, url: str, dest: Path) -> None:
|
||||
env = os.environ.copy()
|
||||
env["DATA_ROOT"] = str(self.data_root)
|
||||
if self.stack_env_path and self.stack_env_path.is_file():
|
||||
for line in self.stack_env_path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
key, _, value = line.partition("=")
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
if key and key not in env:
|
||||
env[key] = value
|
||||
|
||||
cmd = [
|
||||
str(self.download_file_script),
|
||||
"--url",
|
||||
url,
|
||||
"--dest",
|
||||
str(dest),
|
||||
]
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
env=env,
|
||||
cwd=str(self.download_file_script.parent),
|
||||
)
|
||||
job._process = proc
|
||||
assert proc.stdout is not None
|
||||
for line in proc.stdout:
|
||||
job.append_log(line)
|
||||
rc = proc.wait()
|
||||
if job.status == "cancelled":
|
||||
return
|
||||
if rc == 0:
|
||||
job.status = "completed"
|
||||
job.progress = 100
|
||||
else:
|
||||
job.status = "failed"
|
||||
job.error = f"Download script exited with code {rc}"
|
||||
except Exception as exc:
|
||||
job.status = "failed"
|
||||
job.error = str(exc)
|
||||
job.append_log(f"\nERROR: {exc}\n")
|
||||
log.exception("ComfyUI file download failed for %s", job.model_id)
|
||||
finally:
|
||||
job.finished_at = time.time()
|
||||
job._process = None
|
||||
with self._lock:
|
||||
if self._active_by_model.get(job.model_id) == job.id:
|
||||
del self._active_by_model[job.model_id]
|
||||
|
||||
def _run_download(self, job: DownloadJob) -> None:
|
||||
env = os.environ.copy()
|
||||
env["DATA_ROOT"] = str(self.data_root)
|
||||
if self.stack_env_path and self.stack_env_path.is_file():
|
||||
for line in self.stack_env_path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
key, _, value = line.partition("=")
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
if key and key not in env:
|
||||
env[key] = value
|
||||
|
||||
cmd = [str(self.download_script), job.model_id]
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
env=env,
|
||||
cwd=str(self.download_script.parent),
|
||||
)
|
||||
job._process = proc
|
||||
assert proc.stdout is not None
|
||||
for line in proc.stdout:
|
||||
job.append_log(line)
|
||||
rc = proc.wait()
|
||||
if job.status == "cancelled":
|
||||
return
|
||||
if rc == 0:
|
||||
job.status = "completed"
|
||||
job.progress = 100
|
||||
else:
|
||||
job.status = "failed"
|
||||
job.error = f"Download script exited with code {rc}"
|
||||
except Exception as exc:
|
||||
job.status = "failed"
|
||||
job.error = str(exc)
|
||||
job.append_log(f"\nERROR: {exc}\n")
|
||||
log.exception("ComfyUI download failed for %s", job.model_id)
|
||||
finally:
|
||||
job.finished_at = time.time()
|
||||
job._process = None
|
||||
with self._lock:
|
||||
if self._active_by_model.get(job.model_id) == job.id:
|
||||
del self._active_by_model[job.model_id]
|
||||
|
||||
|
||||
def init_download_manager(
|
||||
repo_root: Path,
|
||||
data_root: Path,
|
||||
job_ttl: int = JOB_TTL_SECONDS,
|
||||
) -> DownloadJobManager:
|
||||
global _job_manager
|
||||
ui_root = Path(__file__).resolve().parent
|
||||
catalog_path = comfyui_catalog_path(ui_root)
|
||||
scripts = comfyui_scripts_dir(ui_root)
|
||||
download_script = scripts / "download-model.sh"
|
||||
download_file_script = scripts / "download-file.sh"
|
||||
stack_env = comfyui_stack_env_path(repo_root)
|
||||
_job_manager = DownloadJobManager(
|
||||
catalog_path=catalog_path,
|
||||
data_root=Path(data_root),
|
||||
download_script=download_script,
|
||||
download_file_script=download_file_script,
|
||||
stack_env_path=stack_env if stack_env.is_file() else None,
|
||||
job_ttl=job_ttl,
|
||||
)
|
||||
return _job_manager
|
||||
|
||||
|
||||
def get_download_manager() -> DownloadJobManager:
|
||||
if _job_manager is None:
|
||||
raise ComfyUIModelError("Download manager not initialized")
|
||||
return _job_manager
|
||||
Reference in New Issue
Block a user