feat(v3.12.0): media pipeline, AI image generation, capability probe, companion v2.9.0
Three-month batch sync from internal repo (~80 commits) covering Tracks F.5a, F.7e, F.8, F.17, F.18, F.X. WordPress media pipeline - Pillow-based optimization, AI image generation (OpenAI / Stability / Replicate / Google Nano Banana / OpenRouter), chunked + resumable uploads, bulk delete/reassign, idempotent retries. Capability discovery (F.7e) - Per-site credential probe + adapters for WordPress / WooCommerce / Gitea, tier-fit unions granted ∪ roles, capability badge UI with HTMX partial re-check, install hint in every companion-unreachable error. Companion plugin overhaul - Renamed wordpress-plugin/airano-mcp-seo-bridge → wordpress-plugin/airano-mcp-bridge. - Eight new endpoints: /capabilities, /bulk-meta, /export, /cache-purge, /transient-flush, /site-health, /audit-hook, /upload-and-attach. - wp.org Plugin Check pass: i18n, WP_Filesystem, scheme allowlist on audit-hook URL. Other - Gitea ergonomics (F.17): batch files, tree, search, compare, releases, fork. - Opportunistic bcrypt upgrade for legacy SHA-256 admin keys (F.8). - n8n refactor: structured errors, capability probe, missing tools backfilled. - Idempotency-Key dedup for AI media upload retries; WP client fast-fails on unreachable sites. Docs - README + CLAUDE.md drop the fixed "633 tools" claim. The total grows with each release; per-plugin approximations + dashboard-surfaced counts replace it. - Tools/Tests badges removed in favour of "Plugins: 10". Deployment - PyPI mirror chain, optional BUILD_HTTP_PROXY, Alpine→Yandex apk mirror, Debian-slim Plan-B Dockerfile, mirror.gcr.io variant. CI - Black + Ruff clean on Python 3.12; pytest tests/ green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
204
plugins/ai_image/providers/replicate.py
Normal file
204
plugins/ai_image/providers/replicate.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""Replicate provider (Flux family + other image models).
|
||||
|
||||
Replicate runs predictions asynchronously: POST to ``/v1/predictions``
|
||||
returns a job id, and we poll ``/v1/predictions/{id}`` until status is
|
||||
``succeeded`` or ``failed``. The final output is a URL (or list of URLs)
|
||||
that must be fetched to get bytes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from plugins.ai_image.providers.base import (
|
||||
BaseImageProvider,
|
||||
GenerationRequest,
|
||||
GenerationResult,
|
||||
ProviderError,
|
||||
)
|
||||
|
||||
_logger = logging.getLogger("mcphub.ai_image.replicate")
|
||||
|
||||
_PREDICTIONS_URL = "https://api.replicate.com/v1/predictions"
|
||||
_DEFAULT_MODEL = "black-forest-labs/flux-schnell"
|
||||
_MAX_RETRIES = 3
|
||||
_POLL_INTERVAL = 2.0
|
||||
_POLL_TIMEOUT = 180
|
||||
_REQUEST_TIMEOUT = 60
|
||||
_RETRY_STATUS = {429, 500, 502, 503, 504}
|
||||
|
||||
_COST_TABLE: dict[str, float] = {
|
||||
"black-forest-labs/flux-schnell": 0.003,
|
||||
"black-forest-labs/flux-dev": 0.025,
|
||||
"black-forest-labs/flux-pro": 0.055,
|
||||
}
|
||||
|
||||
|
||||
def _aspect_from_size(size: str) -> str:
|
||||
"""Replicate's Flux models accept an aspect_ratio enum, not WxH."""
|
||||
mapping = {
|
||||
"1024x1024": "1:1",
|
||||
"1344x768": "16:9",
|
||||
"768x1344": "9:16",
|
||||
"1152x896": "3:2",
|
||||
"896x1152": "2:3",
|
||||
}
|
||||
return mapping.get(size, "1:1")
|
||||
|
||||
|
||||
class ReplicateProvider(BaseImageProvider):
|
||||
name = "replicate"
|
||||
|
||||
async def generate(self, api_key: str, request: GenerationRequest) -> GenerationResult:
|
||||
if not api_key:
|
||||
raise ProviderError("PROVIDER_AUTH", "Replicate API token is missing.")
|
||||
|
||||
model = request.model or _DEFAULT_MODEL
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"input": {
|
||||
"prompt": request.prompt,
|
||||
"aspect_ratio": _aspect_from_size(request.size),
|
||||
**(request.extras or {}),
|
||||
},
|
||||
}
|
||||
if request.negative_prompt:
|
||||
payload["input"]["negative_prompt"] = request.negative_prompt
|
||||
|
||||
prediction = await self._create_prediction(api_key, payload)
|
||||
final = await self._poll_until_done(api_key, prediction)
|
||||
|
||||
output = final.get("output")
|
||||
url = _first_url(output)
|
||||
if not url:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
"Replicate prediction finished without an image URL.",
|
||||
{"prediction": final},
|
||||
)
|
||||
image_bytes = await _fetch_url(url)
|
||||
|
||||
meta = {
|
||||
"model": model,
|
||||
"prediction_id": final.get("id"),
|
||||
"status": final.get("status"),
|
||||
"metrics": final.get("metrics", {}),
|
||||
}
|
||||
return GenerationResult(
|
||||
data=image_bytes,
|
||||
mime="image/webp" if url.endswith(".webp") else "image/png",
|
||||
filename=f"replicate-{model.split('/')[-1]}.png",
|
||||
meta=meta,
|
||||
cost_usd=_COST_TABLE.get(model),
|
||||
)
|
||||
|
||||
async def _create_prediction(self, api_key: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Prefer": "wait=0",
|
||||
}
|
||||
delay = 1.0
|
||||
last_error = ""
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
async with session.post(_PREDICTIONS_URL, json=payload, headers=headers) as resp:
|
||||
if resp.status in (200, 201):
|
||||
return await resp.json()
|
||||
text = await resp.text()
|
||||
last_error = text[:500]
|
||||
if resp.status in (401, 403):
|
||||
raise ProviderError(
|
||||
"PROVIDER_AUTH",
|
||||
f"Replicate rejected auth ({resp.status}).",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status == 422:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_REQUEST",
|
||||
f"Replicate rejected request: {last_error}",
|
||||
{"body": last_error},
|
||||
)
|
||||
if resp.status in _RETRY_STATUS and attempt < _MAX_RETRIES:
|
||||
_logger.warning("Replicate %d, retry %d", resp.status, attempt)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
continue
|
||||
if resp.status == 429:
|
||||
raise ProviderError(
|
||||
"PROVIDER_QUOTA",
|
||||
"Replicate rate-limit hit.",
|
||||
{"status": 429, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate HTTP {resp.status}.",
|
||||
{"status": resp.status, "body": last_error},
|
||||
)
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate create_prediction failed after retries: {last_error}",
|
||||
)
|
||||
|
||||
async def _poll_until_done(self, api_key: str, prediction: dict[str, Any]) -> dict[str, Any]:
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
poll_url = (prediction.get("urls") or {}).get("get") or (
|
||||
f"{_PREDICTIONS_URL}/{prediction.get('id')}"
|
||||
)
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
deadline = asyncio.get_event_loop().time() + _POLL_TIMEOUT
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
while True:
|
||||
if asyncio.get_event_loop().time() > deadline:
|
||||
raise ProviderError(
|
||||
"PROVIDER_TIMEOUT",
|
||||
"Replicate prediction did not finish within timeout.",
|
||||
{"prediction_id": prediction.get("id")},
|
||||
)
|
||||
async with session.get(poll_url, headers=headers) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
raise ProviderError(
|
||||
"PROVIDER_UNAVAILABLE",
|
||||
f"Replicate poll HTTP {resp.status}.",
|
||||
{"body": text[:500]},
|
||||
)
|
||||
body = await resp.json()
|
||||
status = body.get("status")
|
||||
if status in ("succeeded",):
|
||||
return body
|
||||
if status in ("failed", "canceled"):
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
f"Replicate prediction {status}: {body.get('error')}",
|
||||
{"prediction": body},
|
||||
)
|
||||
await asyncio.sleep(_POLL_INTERVAL)
|
||||
|
||||
|
||||
def _first_url(output: Any) -> str | None:
|
||||
if isinstance(output, str):
|
||||
return output
|
||||
if isinstance(output, list) and output:
|
||||
first = output[0]
|
||||
if isinstance(first, str):
|
||||
return first
|
||||
return None
|
||||
|
||||
|
||||
async def _fetch_url(url: str) -> bytes:
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url) as resp:
|
||||
if resp.status != 200:
|
||||
raise ProviderError(
|
||||
"PROVIDER_BAD_RESPONSE",
|
||||
f"Failed to fetch Replicate output (HTTP {resp.status}).",
|
||||
{"status": resp.status},
|
||||
)
|
||||
return await resp.read()
|
||||
Reference in New Issue
Block a user