"""Optional features that need packages the pack does not force on anyone. A person switches one on in Settings ▸ Modules; nothing is installed then. The wish is written to addons.json, and the next ComfyUI start installs what is missing (before any module imports it), through the same never-moves-anything-installed path an update uses (update.install_lines). State of an add-on: available (not wanted, not installed), pending (wanted, not installed yet), installed (wanted and installed), disabled (installed, but switched off). Code that depends on one asks `enabled(id)`. An add-on this machine cannot run (no NVIDIA GPU) is not offered at all. """ from __future__ import annotations import importlib import importlib.util import json import os import platform import sys import threading from . import config, log _lock = threading.Lock() _SPACY_MODEL = ("en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/" "en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl") # preshed 3.0.14 breaks spaCy's model loading (OverflowError: Python int too large to convert to C # ssize_t, in Tokenizer.from_disk) on Python 3.13; 3.0.13 loads fine. PRESHED_PIN = "preshed<3.0.14" def _nvidia_gpu() -> bool: """nvidia-vfx ships wheels for x86-64 Linux and Windows only, and needs a CUDA card.""" if sys.platform not in ("linux", "win32") or platform.machine().lower() not in ("x86_64", "amd64"): return False try: import torch return bool(torch.cuda.is_available() and torch.version.cuda and torch.cuda.get_device_capability() >= (7, 5)) # RTX starts at Turing except Exception: # noqa: BLE001 return False ADDONS = ( {"id": "nvidia_vsr", "title": "NVIDIA RTX Super Resolution", "hint": "Native upscaling for images and videos. Installs NVIDIA's own package, under its licence.", "requires": ("nvidia-vfx>=0.2",), "modules": ("nvvfx",), "offered": _nvidia_gpu}, {"id": "nvidia_framegen", "title": "NVIDIA RTX Frame Generation", "hint": "Adds in-between frames to finished videos (x2 to x4). Needs an RTX card and nvidia-vfx 0.2 or newer.", "requires": ("nvidia-vfx>=0.2",), "modules": ("nvvfx",), "floors": {"nvidia-vfx": "0.2"}, "offered": _nvidia_gpu}, {"id": "camera", "title": "Automatic camera and focus", "hint": "Lets Shot camera aim camera moves, cuts, transitions and focus at the right things in your prompts, using spaCy.", "requires": ("spacy>=3.8,<3.9", PRESHED_PIN, _SPACY_MODEL), "modules": ("spacy", "en_core_web_sm"), "offered": lambda: True}, {"id": "lora_training", "title": "LoRA training", "hint": "Train your own LoRA from your media and rated results. Installs onnxruntime for automatic tagging.", "requires": ("onnxruntime>=1.17",), "modules": ("onnxruntime",), "offered": lambda: True}, ) _BY_ID = {a["id"]: a for a in ADDONS} def _version_tuple(v) -> tuple: out = [] for part in str(v).split("."): digits = "".join(ch for ch in part if ch.isdigit()) if not digits: break out.append(int(digits)) return tuple(out) def _old(addon) -> bool: """True when a package the add-on needs is installed below the version it needs ("floors"): present, but not usable.""" from . import update for name, floor in (addon.get("floors") or {}).items(): have = update._installed_version(name) if have is not None and _version_tuple(have) < _version_tuple(floor): return True return False def _present(addon) -> bool: return all(importlib.util.find_spec(m) is not None for m in addon["modules"]) and not _old(addon) def _read() -> dict | None: try: data = json.loads(config.ADDONS_FILE.read_text(encoding="utf-8")) except (OSError, ValueError): return None return data if isinstance(data, dict) else None def _write(data: dict) -> None: tmp = config.ADDONS_FILE.with_suffix(".tmp") try: tmp.write_text(json.dumps(data, indent=1), encoding="utf-8") os.replace(tmp, config.ADDONS_FILE) except OSError as exc: log.warning("add-ons", f"could not write {config.ADDONS_FILE.name}: {exc}. The choice is lost on restart.") def _wanted(data) -> set: """What was asked for. With no file yet, whatever is already installed counts as wanted: the packs people installed before add-ons existed must not go quiet.""" if data is not None and isinstance(data.get("wanted"), list): return {i for i in data["wanted"] if i in _BY_ID} return {a["id"] for a in ADDONS if _present(a)} def enabled(addon_id: str) -> bool: """Wanted and installed: the only state in which code may use the add-on.""" a = _BY_ID.get(addon_id) return bool(a) and a["offered"]() and addon_id in _wanted(_read()) and _present(a) def listing() -> list[dict]: """[{id, title, hint, state, failed}] for every add-on this machine can run.""" data = _read() wanted, failed = _wanted(data), (data or {}).get("failed") or {} out = [] for a in ADDONS: if not a["offered"](): continue here, want = _present(a), a["id"] in wanted out.append({"id": a["id"], "title": a["title"], "hint": a["hint"], "state": ("installed" if here else "available" if failed.get(a["id"]) else "pending") if want else ("disabled" if here else "available"), # a failed install is offered again, not "pending" "failed": failed.get(a["id"]) if not here else None}) return out def set_wanted(addon_id: str, on: bool) -> bool: """Record the wish. False when there is no such add-on here. Turning one on clears its failure, which is how a retry is asked for.""" a = _BY_ID.get(addon_id) if a is None or not a["offered"](): return False with _lock: data = _read() or {} wanted = _wanted(_read()) wanted = wanted | {addon_id} if on else wanted - {addon_id} failed = dict(data.get("failed") or {}) if on: failed.pop(addon_id, None) _write({"wanted": sorted(wanted), "failed": failed}) return True def preshed_broken() -> bool: """True when the installed preshed is the release that cannot load spaCy models.""" try: from importlib import metadata v = tuple(int(x) for x in metadata.version("preshed").split(".")[:3]) except Exception: # noqa: BLE001 return False return v >= (3, 0, 14) def repair_spacy() -> bool: """Downgrade preshed when spaCy is here with the broken one. Only a restart makes the running process see it (the compiled module is already loaded), so the log says so.""" if not (importlib.util.find_spec("spacy") and preshed_broken()): return False import subprocess try: proc = subprocess.run([sys.executable, "-m", "pip", "install", "--disable-pip-version-check", PRESHED_PIN], capture_output=True, text=True, timeout=300) except (OSError, subprocess.SubprocessError) as exc: log.alert("add-ons", f"spaCy needs `pip install \"{PRESHED_PIN}\"` and the automatic try failed: {exc}") return False if proc.returncode != 0: log.alert("add-ons", f"spaCy needs `pip install \"{PRESHED_PIN}\"` and pip failed: {(proc.stderr or proc.stdout)[-300:]}") return False log.info("add-ons", "spaCy could not load its model with preshed 3.0.14; installed an older preshed. Restart ComfyUI once.") return True def install_wanted() -> None: """At ComfyUI start: install what is wanted and missing. Never raises, and never retries a failure on its own (a broken install would add minutes to every start); toggling the add-on retries.""" try: if "camera" in _wanted(_read()): repair_spacy() with _lock: data = _read() or {} failed = dict(data.get("failed") or {}) wanted = _wanted(_read()) todo = [a for a in ADDONS if a["id"] in wanted and a["offered"]() and not _present(a) and a["id"] not in failed] if not todo: return from . import update for a in todo: log.info("add-ons", f"installing {a['title']}…") result = update.install_lines(a["requires"], upgrade=tuple(a.get("floors") or ())) importlib.invalidate_caches() if _present(a): continue failed[a["id"]] = (result.get("detail") or "the install did not finish")[-400:] log.alert("add-ons", f"{a['title']} could not be installed: {failed[a['id']]}") _write({"wanted": sorted(wanted), "failed": failed}) except Exception as exc: # noqa: BLE001 log.failed("core.addons", exc)