diff --git a/src/gpu_rent/civitai.py b/src/gpu_rent/civitai.py index a817bc9..c7e5649 100644 --- a/src/gpu_rent/civitai.py +++ b/src/gpu_rent/civitai.py @@ -80,6 +80,39 @@ def pick_primary_file(version: dict) -> dict | None: return None +def pick_preview_image(version: dict) -> tuple[str, str] | None: + """First usable Civitai preview → (url, sidecar_suffix) for SwarmUI. + + SwarmUI looks for ``stem.preview.jpg`` / ``.preview.png`` / ``.jpg`` etc. + next to the weight (T2IModelHandler AutoImageFormatSuffixes). + """ + images = version.get("images") or [] + if not isinstance(images, list): + return None + for item in images: + if not isinstance(item, dict): + continue + kind = str(item.get("type") or "image").lower() + if kind and kind not in {"image", "img", ""}: + continue + url = str(item.get("url") or "").strip() + if not url.startswith("http"): + continue + path = url.split("?", 1)[0].lower() + if path.endswith(".png"): + suffix = ".preview.png" + elif path.endswith(".jpeg") or path.endswith(".jpg"): + suffix = ".preview.jpg" + elif path.endswith(".webp"): + # SwarmUI auto-format list has no .preview.webp — use .jpg name; + # most Civitai CDN URLs are jpeg without extension. + suffix = ".preview.jpg" + else: + suffix = ".preview.jpg" + return url, suffix + return None + + def fetch_model_version(token: str, host: str, version_id: int, timeout: float = 30.0) -> tuple[str, dict]: """GET /api/v1/model-versions/{id}; one retry on the other Civitai host.""" first = _normalize_host(host) diff --git a/src/gpu_rent/doctor.py b/src/gpu_rent/doctor.py index 9cdba17..a5f0c7b 100644 --- a/src/gpu_rent/doctor.py +++ b/src/gpu_rent/doctor.py @@ -349,17 +349,31 @@ def _local_manifests(cfg: Config, checks: list[Check]) -> None: try: repos = parse_extensions(cfg.extensions_manifest) - if cfg.extensions_manifest.is_file(): - checks.append( - Check( - "extensions.yaml", - True, - True, - f"{len(repos)} git-реп", - ) - ) - else: + if not cfg.extensions_manifest.is_file(): checks.append(Check("extensions.yaml", True, False, "файла нет — стоковый SwarmUI")) + elif repos: + checks.append(Check("extensions.yaml", True, True, f"{len(repos)} git-реп")) + else: + example = cfg.extensions_manifest.with_name("extensions.example.yaml") + n_ex = 0 + if example.is_file(): + try: + n_ex = len(parse_extensions(example)) + except ConfigError: + n_ex = 0 + if n_ex: + checks.append( + Check( + "extensions.yaml", + True, + False, + f"пустой; в example {n_ex} реп — скопируй и gpu-rent seed-extensions", + ) + ) + else: + checks.append( + Check("extensions.yaml", True, False, "0 реп — стоковый SwarmUI") + ) except ConfigError as exc: checks.append(Check("extensions.yaml", False, True, str(exc))) diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index ae7e879..6b980e1 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -11,7 +11,7 @@ from pathlib import Path import httpx -from gpu_rent.civitai import fetch_model_version, pick_primary_file +from gpu_rent.civitai import fetch_model_version, pick_preview_image, pick_primary_file from gpu_rent.config import Config from gpu_rent.errors import CloudError, GpuRentError from gpu_rent.manifests import ( @@ -126,6 +126,14 @@ def seed_extensions(cfg: Config, host: str, log: Log, *, update: bool = True) -> else: log("extensions.yaml пуст — стоковый SwarmUI") return False + if not repos and update: + if not all_repos: + log("extensions.yaml пуст — только update уже установленных на data (если есть)") + else: + log( + "extensions: yaml отфильтрован по requires " + f"(LLM_RUNTIME={runtime}) — только update установленных" + ) jobs = [] for repo in repos: jobs.append( @@ -462,18 +470,22 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None: ), "tags": version.get("tags") or [], } - jobs.append( - { - "dest": dest, - "url": _download_url(api_host, vid, info), - "sha256": sha, - "auth": "civitai", - "sidecars": { - f"{stem}.civitai.json": civitai_json, - f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2), - }, - } - ) + job: dict = { + "dest": dest, + "url": _download_url(api_host, vid, info), + "sha256": sha, + "auth": "civitai", + "sidecars": { + f"{stem}.civitai.json": civitai_json, + f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2), + }, + } + preview = pick_preview_image(version) + if preview: + preview_url, preview_suffix = preview + job["preview_url"] = preview_url + job["preview_dest"] = f"{DATA}/Models/{folder}/{stem}{preview_suffix}" + jobs.append(job) if not jobs: if not cfg.civitai_api_token: diff --git a/src/gpu_rent/remote/civitai_fetch.py b/src/gpu_rent/remote/civitai_fetch.py index 72516f8..1afd18b 100644 --- a/src/gpu_rent/remote/civitai_fetch.py +++ b/src/gpu_rent/remote/civitai_fetch.py @@ -153,6 +153,30 @@ def download(url: str, dest: Path, token: str, *, label: str, auth_host: str) -> partial.replace(dest) +def ensure_preview(job: dict, token: str, *, prefix: str, auth_host: str) -> None: + """Download SwarmUI sidecar preview if missing (even when weight was skipped).""" + preview_url = str(job.get("preview_url") or "").strip() + preview_dest = str(job.get("preview_dest") or "").strip() + if not preview_url or not preview_dest: + return + dest = Path(preview_dest) + if dest.is_file() and dest.stat().st_size > 0: + return + try: + print(f"{prefix} preview: {dest.name}", flush=True) + # Image CDN usually needs no Bearer; still pass token for civitai hosts. + download( + preview_url, + dest, + token if auth_host == "civitai" else "", + label=f"{prefix} {dest.name}", + auth_host=auth_host if auth_host == "civitai" else "civitai", + ) + print(f"{prefix} preview ok {dest.name} ({fmt_bytes(dest.stat().st_size)})") + except Exception as exc: + print(f"{prefix} WARN preview {dest.name}: {exc}") + + def main() -> int: civitai_token = ( TOKEN_PATH.read_text(encoding="utf-8").strip() if TOKEN_PATH.is_file() else "" @@ -187,6 +211,7 @@ def main() -> int: extra = dest.parent / extra_name extra.parent.mkdir(parents=True, exist_ok=True) extra.write_text(extra_text, encoding="utf-8") + ensure_preview(job, token, prefix=prefix, auth_host=auth_host) continue if reason: print(f"{prefix} {reason}: {dest.name}") @@ -212,6 +237,7 @@ def main() -> int: for extra_name, extra_text in (job.get("sidecars") or {}).items(): extra = dest.parent / extra_name extra.write_text(extra_text, encoding="utf-8") + ensure_preview(job, token, prefix=prefix, auth_host=auth_host) print(f"{prefix} ok {dest.name} ({fmt_bytes(dest.stat().st_size)})") except Exception as exc: failed += 1 diff --git a/tests/test_civitai_fetch.py b/tests/test_civitai_fetch.py index 7e91439..1f377ea 100644 --- a/tests/test_civitai_fetch.py +++ b/tests/test_civitai_fetch.py @@ -1,5 +1,6 @@ from pathlib import Path +from gpu_rent.civitai import pick_preview_image from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip @@ -54,3 +55,28 @@ def test_progress_line_unknown_total(): line = progress_line("file", 1024 * 1024, None, 100_000) assert "%" not in line assert "/s" in line + + +def test_pick_preview_image_first_jpeg(): + version = { + "images": [ + {"type": "video", "url": "https://example.com/x.mp4"}, + {"type": "image", "url": "https://image.civitai.com/x/width=450/abc.jpeg"}, + ] + } + got = pick_preview_image(version) + assert got is not None + url, suffix = got + assert "abc.jpeg" in url + assert suffix == ".preview.jpg" + + +def test_pick_preview_image_png_and_default(): + assert pick_preview_image({"images": [{"url": "https://cdn.example/a.png"}]})[1] == ( + ".preview.png" + ) + assert pick_preview_image({"images": [{"url": "https://cdn.example/hash/width=450"}]})[ + 1 + ] == ".preview.jpg" + assert pick_preview_image({"images": []}) is None + assert pick_preview_image({}) is None