tomasz-syn-grzegorza
73e4fc005e
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.
859 lines
28 KiB
Python
859 lines
28 KiB
Python
"""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
|