"""Parse models.yaml / extensions.yaml. version_id 0 is a placeholder.""" from __future__ import annotations import re from dataclasses import dataclass from pathlib import Path from typing import Any from urllib.parse import parse_qs, urlparse import yaml from gpu_rent.errors import ConfigError MODEL_TYPES = ( "checkpoint", "lora", "vae", "embedding", "controlnet", "upscaler", "clip", ) MODEL_DIRS = { "checkpoint": "Stable-Diffusion", "lora": "Lora", "vae": "VAE", "embedding": "Embeddings", "controlnet": "controlnet", "upscaler": "upscale_models", "clip": "clip", } _VERSION_QS = re.compile(r"modelVersionId=(\d+)", re.I) _VERSION_PATH = re.compile(r"/model-versions/(\d+)", re.I) _DOWNLOAD_PATH = re.compile(r"/api/download/models/(\d+)", re.I) _SHA_REF = re.compile(r"^[0-9a-fA-F]{40}$") @dataclass class ModelEntry: kind: str version_id: int | None url: str | None VALID_REQUIRES = frozenset({"none", "ollama", "any-llm"}) @dataclass class GitRepo: kind: str url: str ref: str directory: str | None requires: str = "none" def _load_yaml(path: Path) -> Any: if not path.is_file(): return None text = path.read_text(encoding="utf-8") if not text.strip(): return {} try: return yaml.safe_load(text) except yaml.YAMLError as exc: raise ConfigError(f"Не разобрать YAML {path}: {exc}") from exc def parse_models(path: Path) -> list[ModelEntry]: data = _load_yaml(path) if data is None: return [] if data == {} or data is None: return [] if not isinstance(data, dict): raise ConfigError(f"{path}: корень должен быть mapping типов моделей") entries: list[ModelEntry] = [] for kind, items in data.items(): if kind not in MODEL_TYPES: continue if not items: continue if not isinstance(items, list): raise ConfigError(f"{path}: {kind} должен быть списком") for item in items: if not isinstance(item, dict): raise ConfigError(f"{path}: элемент {kind} — объект с version_id или url") vid = item.get("version_id") url = item.get("url") if vid in (0, "0", None) and not url: continue version_id = int(vid) if vid not in (None, "", 0, "0") else None url_s = str(url) if url else None if version_id is None and url_s: version_id = extract_version_id(url_s) entries.append(ModelEntry(kind=kind, version_id=version_id, url=url_s)) return entries def normalize_requires(value: object | None) -> str: raw = str(value or "none").strip().lower().replace("_", "-") if raw in {"", "none", "always", "any"}: return "none" if raw in VALID_REQUIRES: return raw raise ConfigError( f"requires={value!r}: жду none|ollama|any-llm" ) def repo_matches_runtime(repo: GitRepo, llm_runtime: str) -> bool: """Whether this extension should be cloned for the active LLM_RUNTIME.""" req = normalize_requires(repo.requires) runtime = (llm_runtime or "none").strip().lower() if req == "none": return True if req == "any-llm": return runtime == "ollama" return runtime == req def parse_extensions(path: Path) -> list[GitRepo]: data = _load_yaml(path) if not data: return [] if not isinstance(data, dict): raise ConfigError(f"{path}: корень swarmui: / comfy: / assistent:") repos: list[GitRepo] = [] for kind in ("swarmui", "comfy", "assistent"): items = data.get(kind) or [] if not items: continue if not isinstance(items, list): raise ConfigError(f"{path}: {kind} должен быть списком") for item in items: if not isinstance(item, dict) or not item.get("url"): raise ConfigError(f"{path}: у {kind} нужен url") try: requires = normalize_requires(item.get("requires")) except ConfigError as exc: raise ConfigError(f"{path}: {exc}") from exc repos.append( GitRepo( kind=kind, url=str(item["url"]).strip(), ref=str(item.get("ref") or "main"), directory=str(item["dir"]) if item.get("dir") else None, requires=requires, ) ) return repos def extract_version_id(url: str) -> int | None: text = url.strip() for rx in (_VERSION_QS, _VERSION_PATH, _DOWNLOAD_PATH): match = rx.search(text) if match: return int(match.group(1)) parsed = urlparse(text) ids = parse_qs(parsed.query).get("modelVersionId") or parse_qs(parsed.query).get("modelversionid") if ids: try: return int(ids[0]) except ValueError: return None return None def repo_dirname(repo: GitRepo) -> str: if repo.directory: return repo.directory name = repo.url.rstrip("/").rsplit("/", 1)[-1] if name.endswith(".git"): name = name[:-4] return name or "extension" def is_commit_sha(ref: str) -> bool: return bool(_SHA_REF.match(ref.strip())) def remote_root_for(repo: GitRepo) -> str: if repo.kind == "swarmui": base = "/mnt/swarm_data/Extensions" elif repo.kind == "assistent": base = "/mnt/swarm_data/Assistent/extensions" else: base = "/mnt/swarm_data/DLNodes" return f"{base}/{repo_dirname(repo)}"