Spaces:
Running on Zero
Running on Zero
File size: 3,579 Bytes
4cd8837 4aaae80 4cd8837 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | from __future__ import annotations
import base64
from typing import Any
from hearthnet.bus.capability import CapabilityDescriptor, RouteRequest
from hearthnet.services.image.backends.base import GenerationResult, ImageGenerateBackend
class ImageGenerateService:
"""Service wrapping image-generation backends.
Registers: img.generate@1.0
"""
name = "image.generate"
def __init__(
self,
backends: list[ImageGenerateBackend] | None = None,
bus: Any = None,
) -> None:
self._backends: list[ImageGenerateBackend] = backends if backends is not None else []
self._bus = bus
self._by_name: dict[str, ImageGenerateBackend] = {b.name: b for b in self._backends}
# ── Service registration ──────────────────────────────────────────────────
def capabilities(self) -> list[tuple]:
return [
(
CapabilityDescriptor(
name="img.generate",
max_concurrent=1,
idempotent=False,
timeout_seconds=120,
),
self.generate,
None,
),
]
def register(self, bus: Any) -> None:
self._bus = bus
for cap, handler, predicate in self.capabilities():
bus.register_local(cap, handler, predicate)
# ── Handler ───────────────────────────────────────────────────────────────
async def generate(self, req: RouteRequest) -> dict:
if not self._backends:
return {
"error": "unavailable",
"message": "no image generation backends installed",
}
params: dict = req.body.get("input", {})
prompt: str | None = params.get("prompt")
if not prompt:
return {"error": "bad_request", "message": "prompt required"}
width: int = int(params.get("width", 512))
height: int = int(params.get("height", 512))
steps: int = int(params.get("steps", 20))
lora: str | None = params.get("lora")
backend_name: str | None = params.get("backend")
# Clamp dimensions to sane limits
width = max(64, min(width, 2048))
height = max(64, min(height, 2048))
steps = max(1, min(steps, 200))
# Select backend
backend: ImageGenerateBackend | None = None
if backend_name:
backend = self._by_name.get(backend_name)
if backend is None:
return {"error": "bad_request", "message": f"unknown backend: {backend_name}"}
else:
backend = self._backends[0]
result: GenerationResult = await backend.generate(
prompt, width=width, height=height, steps=steps, lora=lora
)
image_b64 = base64.b64encode(result.image_bytes).decode("ascii")
return {
"output": {
"image_b64": image_b64,
"width": result.width,
"height": result.height,
"backend": result.backend,
"ms": result.ms,
},
"meta": {},
}
def health(self) -> dict:
return {
"service": self.name,
"backends": [b.health() for b in self._backends],
"available": len(self._backends) > 0,
}
|