"""NVIDIA NIM image-generation client for AutoShortsBot keyframes.""" from __future__ import annotations import base64 import json import os import urllib.error import urllib.request from dataclasses import dataclass from pathlib import Path from typing import Any NVIDIA_IMAGE_ENDPOINTS: dict[str, str] = { "qwen-image": "https://ai.api.nvidia.com/v1/genai/qwen/qwen-image", "stable-diffusion-3.5-large": "https://ai.api.nvidia.com/v1/genai/stabilityai/stable-diffusion-3_5-large", # The two Build pages above currently expose this hosted endpoint in their embedded snippets. # Keep it as the production fallback because it is available through NVIDIA API Catalog. "flux.1-dev": "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.1-dev", } @dataclass(frozen=True) class NvidiaImageRequest: """Provider-neutral request for a single keyframe image.""" prompt: str output_path: Path width: int = 768 height: int = 1344 seed: int = 101 steps: int = 28 cfg_scale: float = 4.0 mode: str = "base" @dataclass(frozen=True) class NvidiaImageResult: """Structured result for one provider attempt.""" status: str model: str endpoint: str output_path: Path | None = None http_status: int | None = None error: str | None = None bytes_written: int = 0 class UrllibNvidiaImageTransport: """Small urllib transport so tests can swap the network layer.""" def post_json(self, endpoint: str, api_key: str, payload: dict[str, Any], *, timeout_seconds: int) -> tuple[int, bytes, str]: request = urllib.request.Request( endpoint, data=json.dumps(payload).encode("utf-8"), headers={ "Authorization": f"Bearer {api_key}", "Accept": "application/json", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(request, timeout=timeout_seconds) as response: return response.status, response.read(), response.headers.get("Content-Type", "") except urllib.error.HTTPError as exc: return exc.code, exc.read(), exc.headers.get("Content-Type", "") class NvidiaImageClient: """Client for hosted NVIDIA image-generation endpoints. The Build portal has both hosted API endpoints and deploy/self-host pages. This client only uses the hosted API path under ``https://ai.api.nvidia.com/v1/genai/...`` and does not require hosting a NIM container locally. """ def __init__( self, *, api_key: str | None = None, endpoints: dict[str, str] | None = None, transport: Any | None = None, timeout_seconds: int = 180, ) -> None: self.api_key = api_key or load_nvidia_api_key() self.endpoints = endpoints or NVIDIA_IMAGE_ENDPOINTS self.transport = transport or UrllibNvidiaImageTransport() self.timeout_seconds = timeout_seconds def generate( self, request: NvidiaImageRequest, *, model_order: list[str] | None = None, ) -> NvidiaImageResult: """Generate one image, trying models in order until one succeeds.""" order = model_order or ["qwen-image", "stable-diffusion-3.5-large", "flux.1-dev"] last_result: NvidiaImageResult | None = None for model in order: endpoint = self.endpoints[model] result = self._generate_with_model(request, model=model, endpoint=endpoint) if result.status == "success": return result last_result = result if last_result is None: raise ValueError("model_order is empty") return last_result def _generate_with_model(self, request: NvidiaImageRequest, *, model: str, endpoint: str) -> NvidiaImageResult: payload = { "prompt": request.prompt, "mode": request.mode, "cfg_scale": request.cfg_scale, "width": request.width, "height": request.height, "seed": request.seed, "steps": request.steps, } status_code, body, _content_type = self.transport.post_json( endpoint, self.api_key, payload, timeout_seconds=self.timeout_seconds, ) if status_code < 200 or status_code >= 300: return NvidiaImageResult( status="failed", model=model, endpoint=endpoint, http_status=status_code, error=_safe_error_text(body), ) try: image_bytes = _extract_image_bytes(json.loads(body.decode("utf-8"))) except Exception as exc: # noqa: BLE001 - convert provider quirks into reportable failures. return NvidiaImageResult( status="failed", model=model, endpoint=endpoint, http_status=status_code, error=f"could not parse image response: {exc}", ) request.output_path.parent.mkdir(parents=True, exist_ok=True) request.output_path.write_bytes(image_bytes) return NvidiaImageResult( status="success", model=model, endpoint=endpoint, output_path=request.output_path, http_status=status_code, bytes_written=len(image_bytes), ) def load_nvidia_api_key(path: Path | None = None) -> str: """Load NVIDIA API key from env or Hermes secret file without printing it.""" for env_name in ("NVIDIA_API_KEY", "NGC_API_KEY", "NVIDIA_NIM_API_KEY"): value = os.getenv(env_name) if value: return value.strip() secret_path = path or Path.home() / ".hermes" / "secrets" / "nvidia_api_key" raw = secret_path.read_text(encoding="utf-8").strip() if "=" in raw: name, value = raw.split("=", 1) if name.strip() not in {"NVIDIA_API_KEY", "NGC_API_KEY", "NVIDIA_NIM_API_KEY"}: raise ValueError(f"Unsupported NVIDIA secret key name: {name.strip()}") return value.strip() return raw.strip() def _extract_image_bytes(payload: dict[str, Any]) -> bytes: """Extract image bytes from common NVIDIA Visual GenAI response shapes.""" if artifacts := payload.get("artifacts"): first = artifacts[0] if encoded := first.get("base64"): return base64.b64decode(encoded) for key in ("image", "b64_json", "base64"): encoded = payload.get(key) if isinstance(encoded, str) and encoded: return base64.b64decode(encoded) if data := payload.get("data"): first = data[0] for key in ("b64_json", "base64", "image"): encoded = first.get(key) if isinstance(encoded, str) and encoded: return base64.b64decode(encoded) raise ValueError(f"no base64 image field found; keys={sorted(payload.keys())}") def _safe_error_text(body: bytes, *, limit: int = 500) -> str: text = body.decode("utf-8", errors="replace").strip() return text[:limit]