Enhance SwarmUI integration and GPU environment verification
- Updated CLI documentation to reflect the new handling of `CIVITAI_API_TOKEN`, which is now automatically passed to SwarmUI user settings during startup. - Improved the `render_access_panel` function to include additional warnings for idle-killer failures and stack errors, enhancing user feedback. - Introduced a new function `seed_swarmui_api_keys` to manage API key injection into SwarmUI, ensuring seamless integration with the Model Downloader. - Enhanced GPU environment verification logic to include fail-fast checks for critical components like CUDA, improving error handling and user notifications. - Updated tests to validate the new API key handling and access panel behavior, ensuring robustness in the integration process.
This commit is contained in:
@@ -0,0 +1,90 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Push upstream API keys into SwarmUI (user GenericData via SetAPIKey).
|
||||
|
||||
Stdlib only. Keys file: /tmp/gpu-rent-swarm-api-keys.json (mode 600), shape:
|
||||
{"civitai_api": "...", "huggingface_api": "..."} # omit empty
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
KEYS_PATH = Path("/tmp/gpu-rent-swarm-api-keys.json")
|
||||
SWARM = "http://127.0.0.1:7801"
|
||||
ACCEPTED = ("civitai_api", "huggingface_api", "stability_api")
|
||||
|
||||
|
||||
def post(path: str, payload: dict, timeout: float = 15.0) -> dict:
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
f"{SWARM}{path}",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
|
||||
|
||||
def wait_session(deadline: float) -> str:
|
||||
last = ""
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
data = post("/API/GetNewSession", {})
|
||||
sid = str(data.get("session_id") or "")
|
||||
if sid:
|
||||
return sid
|
||||
last = "no session_id"
|
||||
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
|
||||
last = str(exc)[:160]
|
||||
time.sleep(2)
|
||||
raise SystemExit(f"SwarmUI session unavailable: {last}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if not KEYS_PATH.is_file():
|
||||
print("no keys file — skip")
|
||||
return 0
|
||||
try:
|
||||
raw = json.loads(KEYS_PATH.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
print(f"bad keys file: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
finally:
|
||||
try:
|
||||
KEYS_PATH.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
keys = {k: str(v).strip() for k, v in (raw or {}).items() if k in ACCEPTED and str(v).strip()}
|
||||
if not keys:
|
||||
print("no api keys to set")
|
||||
return 0
|
||||
|
||||
sid = wait_session(time.time() + 120)
|
||||
for key_type, value in keys.items():
|
||||
try:
|
||||
resp = post(
|
||||
"/API/SetAPIKey",
|
||||
{"session_id": sid, "keyType": key_type, "key": value},
|
||||
)
|
||||
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
|
||||
print(f"SetAPIKey {key_type} failed: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
if resp.get("error"):
|
||||
print(f"SetAPIKey {key_type}: {resp['error']}", file=sys.stderr)
|
||||
return 1
|
||||
if not resp.get("success"):
|
||||
print(f"SetAPIKey {key_type}: unexpected {resp}", file=sys.stderr)
|
||||
return 1
|
||||
print(f"SetAPIKey {key_type}=ok")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user