#!/usr/bin/env python3
"""
Krea-2 as-a-Service — HTTP API wrapper
Espone il client krea/Krea-2 come REST API GET, ideale per agent AI / tool calling.
Uso:
python api.py # avvia server su :7860
python api.py --port 8080 --host 0.0.0.0
Endpoints principali (tutti GET):
GET / → info + docs
GET /health → healthcheck
GET /generate?prompt=... → genera 1 immagine (sync)
GET /generate?prompt=...&n=5 → genera N in parallelo
GET /batch?prompts=a|b|c → batch da lista pipe-separated
GET /jobs/{job_id} → stato job async
GET /jobs/{job_id}/result → risultato quando pronto
GET /image/{filename} → scarica file generato
GET /stats → statistiche DB proxy
GET /openapi.json → schema OpenAPI per agent
pip install curl_cffi fastapi uvicorn
"""
import argparse
import asyncio
import json
import os
import re
import sys
import time
import uuid
import random
import string
import sqlite3
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional, Literal
from datetime import datetime
from curl_cffi import requests as cffi_requests
from fastapi import FastAPI, HTTPException, Query, BackgroundTasks, Request
from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
# ═══════════════════════════════════════════════════════════════
# CONFIG
# ═══════════════════════════════════════════════════════════════
HF_BASE = "https://huggingface.co"
SPACE_BASE = "https://krea-krea-2.hf.space"
SPACE_ID = "krea/Krea-2"
JWT_URL = f"{HF_BASE}/api/spaces/{SPACE_ID}/jwt"
JOIN_URL = f"{SPACE_BASE}/gradio_api/queue/join"
DATA_URL = f"{SPACE_BASE}/gradio_api/queue/data"
FN_RESOLUTION = 3
FN_GENERATE = 4
DEFAULT_WORKERS = 10
WAVE_SIZE = 25
WAVE_DELAY_MS = 150
TOKEN_TIMEOUT = 8
STREAM_TIMEOUT = 90
DB_PATH = Path("proxy_pool.db")
IMAGES_DIR = Path("generated_images")
IMAGES_DIR.mkdir(exist_ok=True)
BAN_TTL_SECONDS = 3600 * 6
GOOD_TTL_DAYS = 7
MIN_GOOD_RATIO = 0.4
MAX_PARALLEL_JOBS = 10
# Job registry TTL (dopo quanto pulire i job vecchi)
JOB_TTL_SECONDS = 3600 # 1h
PROXYSCRAPE_URLS = [
"https://api.proxyscrape.com/v2/?request=displayproxies&protocol=http&timeout=3000&country=all&ssl=all&anonymity=all",
"https://api.proxyscrape.com/v2/?request=displayproxies&protocol=socks4&timeout=3000&country=all&ssl=all&anonymity=all",
"https://api.proxyscrape.com/v2/?request=displayproxies&protocol=socks5&timeout=3000&country=all&ssl=all&anonymity=all",
]
IMPERSONATE_POOL = [
("chrome124", "Windows NT 10.0; Win64; x64", "Chrome/124.0.0.0"),
("chrome123", "Macintosh; Intel Mac OS X 10_15_7", "Chrome/123.0.0.0"),
("chrome120", "X11; Linux x86_64", "Chrome/120.0.0.0"),
("chrome116", "Windows NT 10.0; Win64; x64", "Chrome/116.0.0.0"),
]
LANGS = [
"en-US,en;q=0.9", "it-IT,it;q=0.9,en;q=0.8",
"en-GB,en;q=0.9", "fr-FR,fr;q=0.9,en;q=0.8",
"de-DE,de;q=0.9,en;q=0.8", "es-ES,es;q=0.9,en;q=0.8",
]
RESOLUTIONS = {
"square": (1024, 1024),
"portrait": (1024, 1536),
"landscape": (1536, 1024),
"square2k": (2048, 2048),
}
# ═══════════════════════════════════════════════════════════════
# PROXY DB (SQLite)
# ═══════════════════════════════════════════════════════════════
class ProxyDB:
def __init__(self, path=DB_PATH):
self.path = path
self._lock = threading.Lock()
self._init_db()
def _conn(self):
conn = sqlite3.connect(self.path, timeout=10, check_same_thread=False)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA synchronous=NORMAL")
return conn
def _init_db(self):
with self._conn() as c:
c.execute("""
CREATE TABLE IF NOT EXISTS proxies (
url TEXT PRIMARY KEY, protocol TEXT NOT NULL,
successes INTEGER DEFAULT 0, failures INTEGER DEFAULT 0,
avg_latency_ms REAL DEFAULT 0.0,
last_success REAL DEFAULT 0, last_failure REAL DEFAULT 0,
banned_until REAL DEFAULT 0,
first_seen REAL DEFAULT 0, score REAL DEFAULT 0
)""")
c.execute("CREATE INDEX IF NOT EXISTS idx_score ON proxies(score DESC)")
c.commit()
def upsert_many(self, proxies):
with self._lock, self._conn() as c:
now = time.time()
for p in proxies:
c.execute("INSERT OR IGNORE INTO proxies (url, protocol, first_seen) VALUES (?, ?, ?)",
(p["url"], p["protocol"], now))
c.commit()
def get_ranked(self, limit, min_good_ratio=0.4, exclude: set = None):
exclude = exclude or set()
with self._lock, self._conn() as c:
now = time.time()
n_good = int(limit * min_good_ratio)
n_new = limit - n_good
good = c.execute("""
SELECT url, protocol, successes, failures, avg_latency_ms, score FROM proxies
WHERE banned_until < ? AND successes > 0
ORDER BY score DESC, avg_latency_ms ASC LIMIT ?
""", (now, n_good * 2)).fetchall()
new_ones = c.execute("""
SELECT url, protocol, successes, failures, avg_latency_ms, score FROM proxies
WHERE banned_until < ? AND successes = 0 AND failures < 3
ORDER BY RANDOM() LIMIT ?
""", (now, n_new * 2)).fetchall()
rows = [r for r in list(good) + list(new_ones) if r[0] not in exclude][:limit]
return [{"url": r[0], "protocol": r[1], "successes": r[2],
"failures": r[3], "avg_latency_ms": r[4], "score": r[5]}
for r in rows]
def mark_success(self, url, latency_ms):
with self._lock, self._conn() as c:
row = c.execute("SELECT successes, avg_latency_ms FROM proxies WHERE url=?", (url,)).fetchone()
if row:
s, old_lat = row
new_s = s + 1
new_lat = old_lat * 0.7 + latency_ms * 0.3 if old_lat > 0 else latency_ms
score = new_s * 100 - (new_lat / 100)
c.execute("UPDATE proxies SET successes=?, avg_latency_ms=?, last_success=?, banned_until=0, score=? WHERE url=?",
(new_s, new_lat, time.time(), score, url))
c.commit()
def mark_failure(self, url, ban=False):
with self._lock, self._conn() as c:
row = c.execute("SELECT successes, failures FROM proxies WHERE url=?", (url,)).fetchone()
if not row: return
s, f = row
new_f = f + 1
banned_until = 0
if ban or (new_f >= 3 and s == 0):
banned_until = time.time() + BAN_TTL_SECONDS
score = s * 100 - new_f * 10
c.execute("UPDATE proxies SET failures=?, last_failure=?, banned_until=?, score=? WHERE url=?",
(new_f, time.time(), banned_until, score, url))
c.commit()
def stats(self):
with self._lock, self._conn() as c:
r = c.execute("""SELECT COUNT(*), SUM(CASE WHEN successes>0 THEN 1 ELSE 0 END),
SUM(CASE WHEN banned_until>? THEN 1 ELSE 0 END),
SUM(CASE WHEN successes=0 AND failures=0 THEN 1 ELSE 0 END)
FROM proxies""", (time.time(),)).fetchone()
return {"total": r[0] or 0, "good": r[1] or 0,
"banned": r[2] or 0, "untested": r[3] or 0}
DB = ProxyDB()
def _fetch_one(url):
try:
r = cffi_requests.get(url, timeout=8, impersonate="chrome120")
if r.status_code == 200:
protocol = "http"
if "socks4" in url: protocol = "socks4"
elif "socks5" in url: protocol = "socks5"
out = []
for line in r.text.strip().split("\n"):
line = line.strip()
if ":" not in line: continue
ip_port = line.split()[0]
proxy_url = (f"{protocol}://{ip_port}" if protocol != "http"
else f"http://{ip_port}")
out.append({"url": proxy_url, "protocol": protocol})
return out
except Exception: pass
return []
def refresh_proxy_pool():
all_proxies = []
with ThreadPoolExecutor(max_workers=3) as pool:
for chunk in pool.map(_fetch_one, PROXYSCRAPE_URLS):
all_proxies.extend(chunk)
DB.upsert_many(all_proxies)
return len(all_proxies)
# ═══════════════════════════════════════════════════════════════
# UTILS
# ═══════════════════════════════════════════════════════════════
def _rand_hash(n=11):
return "".join(random.choices(string.ascii_lowercase + string.digits, k=n))
def _new_identity():
imp, os_str, ver = random.choice(IMPERSONATE_POOL)
ua = f"Mozilla/5.0 ({os_str}) AppleWebKit/537.36 (KHTML, like Gecko) {ver} Safari/537.36"
return {"impersonate": imp, "ua": ua,
"accept_lang": random.choice(LANGS),
"zerogpu_uuid": str(uuid.uuid4())}
def _resolution_label(w, h):
if w == h and w >= 2048: return "Square · 2K"
if w == h: return "Square · 1024"
if w > h: return "Landscape · 1024"
return "Portrait · 1024"
# ═══════════════════════════════════════════════════════════════
# JOB STATE (identico al client, senza UI)
# ═══════════════════════════════════════════════════════════════
@dataclass
class WorkerInfo:
id: int
proxy: str = ""
status: str = "idle"
step: int = 0
total_steps: int = 8
started_at: float = 0.0
@dataclass
class GenJob:
job_id: str
prompt: str
negative_prompt: str = ""
model: str = "Turbo"
steps: int = 8
guidance: float = 0.0
width: int = 1024
height: int = 1024
seed: int = 0
randomize: bool = True
status: str = "pending" # pending|running|success|failed
created_at: float = 0.0
started_at: float = 0.0
completed_at: float = 0.0
# risultati
filename: str = "" # nome file generato
file_path: str = "" # path assoluto
file_size: int = 0
hf_url: str = "" # URL originale su HF
local_url: str = "" # URL da chiamare per scaricare
result_seed: int = 0
error: str = ""
# dettagli race
winner_proxy: str = ""
winner_latency_ms: float = 0.0
workers_launched: int = 0
workers_dead: int = 0
tokens_ok: int = 0
# runtime (non serializzati in API)
workers: dict = field(default_factory=dict)
winner_event: threading.Event = field(default_factory=threading.Event)
lock: threading.Lock = field(default_factory=threading.Lock)
def to_public_dict(self) -> dict:
d = {
"job_id": self.job_id,
"status": self.status,
"prompt": self.prompt,
"negative_prompt": self.negative_prompt,
"params": {
"model": self.model, "steps": self.steps,
"guidance": self.guidance,
"width": self.width, "height": self.height,
"seed": self.seed, "randomize": self.randomize,
},
"timing": {
"created_at": self.created_at,
"started_at": self.started_at or None,
"completed_at": self.completed_at or None,
"duration_s": (self.completed_at - self.started_at) if self.completed_at else None,
"queue_wait_s": (self.started_at - self.created_at) if self.started_at else None,
},
"workers_stats": {
"launched": self.workers_launched,
"dead": self.workers_dead,
"tokens_ok": self.tokens_ok,
},
}
if self.status == "success":
d["result"] = {
"filename": self.filename,
"file_size_bytes": self.file_size,
"seed": self.result_seed,
"hf_url": self.hf_url,
"local_url": self.local_url,
"winner_proxy": self.winner_proxy,
"winner_latency_ms": self.winner_latency_ms,
}
elif self.status == "failed":
d["error"] = self.error
return d
# Registry globale dei job
JOBS: dict[str, GenJob] = {}
JOBS_LOCK = threading.Lock()
USED_PROXIES: set = set()
USED_PROXIES_LOCK = threading.Lock()
def cleanup_old_jobs():
"""Rimuove job più vecchi di JOB_TTL_SECONDS."""
now = time.time()
with JOBS_LOCK:
to_del = [jid for jid, j in JOBS.items()
if (j.completed_at or j.created_at) < now - JOB_TTL_SECONDS]
for jid in to_del:
del JOBS[jid]
return len(to_del)
# ═══════════════════════════════════════════════════════════════
# WORKER (invariato dal client v9)
# ═══════════════════════════════════════════════════════════════
def worker_thread(job: GenJob, worker_id: int, proxy: dict):
st = job.workers[worker_id]
st.proxy = proxy["url"]
st.total_steps = job.steps
if job.winner_event.is_set():
st.status = "killed"; return None
t_start = time.time()
ident = _new_identity()
proxies = {"http": proxy["url"], "https": proxy["url"]}
st.status = "connecting"
try:
s = cffi_requests.Session(impersonate=ident["impersonate"], proxies=proxies)
s.headers.update({"user-agent": ident["ua"], "accept-language": ident["accept_lang"]})
except Exception:
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"], ban=True)
return None
try:
# TOKEN
try:
r = s.get(JWT_URL, headers={
"accept": "*/*", "referer": f"{HF_BASE}/spaces/{SPACE_ID}",
"origin": HF_BASE}, timeout=TOKEN_TIMEOUT)
if r.status_code != 200: raise Exception()
token = (r.json().get("token") or r.json().get("jwt"))
if not token: raise Exception()
except Exception:
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"])
return None
if job.winner_event.is_set(): st.status = "killed"; return None
st.status = "token"
with job.lock: job.tokens_ok += 1
auth = {"accept": "*/*", "content-type": "application/json",
"origin": SPACE_BASE, "referer": f"{SPACE_BASE}/?__theme=system",
"x-gradio-server": f"{SPACE_BASE}/", "x-gradio-user": "app",
"x-zerogpu-token": token, "x-zerogpu-uuid": ident["zerogpu_uuid"]}
# RESOLUTION
sh1 = _rand_hash()
try:
s.post(f"{JOIN_URL}?__theme=system", json={
"data": [_resolution_label(job.width, job.height)],
"fn_index": FN_RESOLUTION, "trigger_id": 10, "session_hash": sh1,
}, headers=auth, timeout=8)
except Exception: pass
if job.winner_event.is_set(): st.status = "killed"; return None
# GENERATE
sh2 = _rand_hash()
try:
r = s.post(f"{JOIN_URL}?__theme=system", json={
"data": [job.prompt, job.negative_prompt or None, job.model,
job.steps, job.guidance, job.width, job.height,
job.seed, job.randomize],
"fn_index": FN_GENERATE, "trigger_id": 7, "session_hash": sh2,
}, headers=auth, timeout=10)
if r.status_code != 200:
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"])
return None
except Exception:
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"])
return None
st.status = "queued"; st.started_at = time.time()
if job.winner_event.is_set(): st.status = "killed"; return None
# STREAM
try:
stream_r = s.get(f"{DATA_URL}?session_hash={sh2}", headers={
"accept": "text/event-stream",
"referer": f"{SPACE_BASE}/?__theme=system",
"x-gradio-server": f"{SPACE_BASE}/"}, stream=True, timeout=STREAM_TIMEOUT)
if stream_r.status_code != 200:
st.status = "dead"; DB.mark_failure(proxy["url"]); return None
except Exception:
st.status = "dead"; DB.mark_failure(proxy["url"]); return None
st.status = "streaming"
for raw in stream_r.iter_lines():
if job.winner_event.is_set(): st.status = "killed"; return None
if not raw: continue
if isinstance(raw, bytes): raw = raw.decode("utf-8", errors="ignore")
if not raw.startswith("data:"): continue
body = raw[5:].strip()
if not body: continue
try: evt = json.loads(body)
except: continue
msg = evt.get("msg")
if msg == "progress":
pd = evt.get("progress_data") or []
if pd and pd[0].get("index") is not None:
st.step = pd[0]["index"] + 1
st.total_steps = pd[0]["length"]
elif msg == "process_completed":
out = evt.get("output", {})
if out.get("error") or not out.get("data"):
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"])
return None
latency = (time.time() - t_start) * 1000
DB.mark_success(proxy["url"], latency)
with job.lock:
if not job.winner_event.is_set():
job.winner_event.set()
job.winner_proxy = proxy["url"]
job.winner_latency_ms = latency
st.status = "done"
return out
return None
elif msg in ("unexpected_error", "close_stream"):
st.status = "dead"; DB.mark_failure(proxy["url"])
return None
st.status = "dead"
DB.mark_failure(proxy["url"])
return None
except Exception:
st.status = "dead"
with job.lock: job.workers_dead += 1
DB.mark_failure(proxy["url"])
return None
finally:
try: s.close()
except: pass
def execute_job(job: GenJob, n_workers: int = DEFAULT_WORKERS,
base_url: str = ""):
"""Esegue un job completo. Chiamata sync bloccante."""
job.status = "running"
job.started_at = time.time()
# Prendi proxy escludendo quelli usati da altri job in corso
with USED_PROXIES_LOCK:
proxies = DB.get_ranked(limit=n_workers,
min_good_ratio=MIN_GOOD_RATIO,
exclude=USED_PROXIES)
for p in proxies:
USED_PROXIES.add(p["url"])
if not proxies:
job.status = "failed"
job.error = "No proxies available in pool. Try /admin/refresh"
job.completed_at = time.time()
return
n_workers = min(n_workers, len(proxies))
job.workers_launched = n_workers
for i in range(n_workers):
job.workers[i] = WorkerInfo(id=i, total_steps=job.steps)
try:
with ThreadPoolExecutor(max_workers=n_workers) as pool:
futures = []
for wave_start in range(0, n_workers, WAVE_SIZE):
if job.winner_event.is_set(): break
wave_end = min(wave_start + WAVE_SIZE, n_workers)
for i in range(wave_start, wave_end):
futures.append(pool.submit(worker_thread, job, i, proxies[i]))
time.sleep(WAVE_DELAY_MS / 1000)
for f in as_completed(futures):
result = f.result()
if result and result.get("data"):
data = result["data"]
img = data[0]
job.result_seed = data[1] if len(data) > 1 else 0
url = img.get("url") if isinstance(img, dict) else img
if not url.startswith("http"):
url = f"{SPACE_BASE}/gradio_api/file={url}"
job.hf_url = url
# Download nel filesystem locale
try:
filename = f"{job.job_id}.png"
file_path = IMAGES_DIR / filename
ident = _new_identity()
with cffi_requests.Session(impersonate=ident["impersonate"]) as s:
s.headers["user-agent"] = ident["ua"]
r = s.get(url, timeout=60)
r.raise_for_status()
file_path.write_bytes(r.content)
job.filename = filename
job.file_path = str(file_path.absolute())
job.file_size = len(r.content)
job.local_url = f"{base_url}/image/{filename}"
job.status = "success"
except Exception as e:
job.status = "failed"
job.error = f"download failed: {e}"
break
finally:
with USED_PROXIES_LOCK:
for p in proxies:
USED_PROXIES.discard(p["url"])
if job.status == "running":
job.status = "failed"
job.error = "No worker succeeded (all proxies failed or quota exceeded)"
job.completed_at = time.time()
def build_job(prompt: str, **kwargs) -> GenJob:
"""Factory di GenJob con validazioni."""
# Risoluzione: preset o custom
width = kwargs.get("width") or 1024
height = kwargs.get("height") or 1024
if kwargs.get("resolution"):
preset = kwargs["resolution"]
if preset not in RESOLUTIONS:
raise ValueError(f"resolution must be one of {list(RESOLUTIONS.keys())}")
width, height = RESOLUTIONS[preset]
# Validazione parametri (da schema /gradio_api/info)
model = kwargs.get("model", "Turbo")
if model not in ("Turbo", "Raw"):
raise ValueError("model must be 'Turbo' or 'Raw'")
steps = int(kwargs.get("steps", 8))
if not 1 <= steps <= 50:
raise ValueError("steps must be between 1 and 50")
guidance = float(kwargs.get("guidance", 0.0))
if not 0.0 <= guidance <= 10.0:
raise ValueError("guidance must be between 0.0 and 10.0")
if not 512 <= width <= 2048:
raise ValueError("width must be between 512 and 2048")
if not 512 <= height <= 2048:
raise ValueError("height must be between 512 and 2048")
seed = kwargs.get("seed")
randomize = seed is None or seed == 0
if not randomize:
seed = int(seed)
if not 0 <= seed <= 2147483647:
raise ValueError("seed must be between 0 and 2147483647")
else:
seed = random.randint(0, 2**31 - 1)
job_id = f"job_{int(time.time()*1000)}_{_rand_hash(6)}"
job = GenJob(
job_id=job_id,
prompt=prompt,
negative_prompt=kwargs.get("negative_prompt", ""),
model=model, steps=steps, guidance=guidance,
width=width, height=height,
seed=seed, randomize=randomize,
created_at=time.time(),
)
with JOBS_LOCK:
JOBS[job_id] = job
return job
# ═══════════════════════════════════════════════════════════════
# FASTAPI APP
# ═══════════════════════════════════════════════════════════════
app = FastAPI(
title="Krea-2 API",
description=(
"REST API for AI image generation via krea/Krea-2 HuggingFace Space.\n"
"Optimized for agentic AI tool-calling: all endpoints are GET, "
"structured JSON responses, OpenAPI schema available at /openapi.json"
),
version="1.0.0",
docs_url="/docs",
redoc_url="/redoc",
openapi_url="/openapi.json",
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
def get_base_url(request: Request) -> str:
"""Costruisce l'URL base per link assoluti."""
scheme = request.headers.get("x-forwarded-proto", request.url.scheme)
host = request.headers.get("host", request.url.hostname)
return f"{scheme}://{host}"
# ─── HOME / INFO ──────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
def home(request: Request):
"""Landing page con documentazione HTML."""
base = get_base_url(request)
stats = DB.stats()
n_jobs = len(JOBS)
n_running = sum(1 for j in JOBS.values() if j.status == "running")
return f"""
Krea-2 API
⚡ Krea-2 API
AI image generation as a service — anonymous, proxy-load-balanced, agent-ready
{stats['total']}proxies
{stats['good']}good
{n_jobs}total jobs
{n_running}running
🔌 Quick Start
GET /generate?prompt=a+cyberpunk+cat
SYNC
Blocca fino al completamento (~5-30s), ritorna JSON con URL immagine.
GET /generate?prompt=...&async=true
ASYNC
Ritorna subito job_id, poi polling con /jobs/{{job_id}}.
GET /generate?prompt=...&n=5
BATCH
Genera N varianti in parallelo (max 10).
📋 All Endpoints
GET / → this page
GET /health → health check
GET /docs → Swagger UI
GET /openapi.json → OpenAPI schema (for agents)
GET /generate → generate image(s)
GET /batch → batch from prompts list
GET /jobs → list all jobs
GET /jobs/{{job_id}} → job status
GET /jobs/{{job_id}}/result → job result (waits if running)
GET /image/{{filename}} → download generated PNG
GET /stats → proxy pool stats
GET /admin/refresh → refresh proxy pool
GET /admin/cleanup → cleanup old jobs
🎯 Example — Agent tool call
curl "{base}/generate?prompt=sunset&model=Turbo&steps=8"
{{
"job_id": "job_1734567890123_abc123",
"status": "success",
"prompt": "sunset",
"result": {{
"filename": "job_1734567890123_abc123.png",
"local_url": "{base}/image/job_1734567890123_abc123.png",
"hf_url": "https://krea-krea-2.hf.space/...",
"seed": 823641299
}}
}}
Swagger docs · OpenAPI · Stats
"""
@app.get("/health")
def health():
"""Health check per load balancer / orchestrator."""
return {
"status": "ok",
"service": "krea-2-api",
"version": "1.0.0",
"timestamp": datetime.utcnow().isoformat() + "Z",
"uptime_s": time.time() - START_TIME,
}
# ─── GENERATE (main endpoint) ──────────────────────────────────
@app.get("/generate")
def generate(
request: Request,
background_tasks: BackgroundTasks,
prompt: str = Query(..., description="Text prompt to generate image from",
min_length=1, max_length=2000, example="a cyberpunk cat"),
negative_prompt: str = Query("", description="What to avoid in generation"),
model: Literal["Turbo", "Raw"] = Query("Turbo",
description="Turbo = fast (8 steps), Raw = quality (more steps)"),
steps: int = Query(8, ge=1, le=50,
description="Denoising steps (Turbo: 4-8, Raw: 20-50)"),
guidance: float = Query(0.0, ge=0.0, le=10.0,
description="Classifier-free guidance scale"),
width: int = Query(1024, ge=512, le=2048),
height: int = Query(1024, ge=512, le=2048),
resolution: Optional[Literal["square","portrait","landscape","square2k"]] = Query(
None, description="Preset resolution (overrides width/height)"),
seed: int = Query(0, ge=0, le=2147483647,
description="Random seed (0 = random for each generation)"),
n: int = Query(1, ge=1, le=MAX_PARALLEL_JOBS,
description=f"Number of images to generate in parallel (max {MAX_PARALLEL_JOBS})"),
async_mode: bool = Query(False, alias="async",
description="If true, return job_id immediately; poll /jobs/{id} for result"),
workers: int = Query(DEFAULT_WORKERS, ge=10, le=200,
description="Proxy race workers per job (higher = faster but heavier)"),
):
"""
**Generate 1 or more AI images from a text prompt.**
Ideal for agentic AI tools. Two modes:
- `async=false` (default): wait for completion, return full result
- `async=true`: return `job_id`, poll `/jobs/{job_id}` for status
Batch: use `n=5` to generate 5 variants of the same prompt in parallel.
"""
base_url = get_base_url(request)
try:
jobs = []
for i in range(n):
job = build_job(
prompt=prompt, negative_prompt=negative_prompt,
model=model, steps=steps, guidance=guidance,
width=width, height=height,
resolution=resolution,
seed=(seed if seed > 0 else None),
)
jobs.append(job)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
if async_mode:
# Lancia in background, ritorna subito i job_ids
for job in jobs:
background_tasks.add_task(execute_job, job, workers, base_url)
return JSONResponse({
"status": "queued",
"job_ids": [j.job_id for j in jobs],
"count": len(jobs),
"poll_urls": [f"{base_url}/jobs/{j.job_id}" for j in jobs],
"message": f"Poll GET /jobs/{{job_id}} to check status",
})
# Sync mode: aspetta tutti
if len(jobs) == 1:
execute_job(jobs[0], workers, base_url)
job = jobs[0]
if job.status == "success":
return job.to_public_dict()
raise HTTPException(status_code=500, detail=job.to_public_dict())
# Batch sync
with ThreadPoolExecutor(max_workers=n) as pool:
futs = {pool.submit(execute_job, j, workers, base_url): j for j in jobs}
for f in as_completed(futs): pass
results = [j.to_public_dict() for j in jobs]
n_ok = sum(1 for j in jobs if j.status == "success")
return {
"status": "batch_complete",
"total": len(jobs),
"success": n_ok,
"failed": len(jobs) - n_ok,
"results": results,
}
@app.get("/batch")
def batch(
request: Request,
background_tasks: BackgroundTasks,
prompts: str = Query(...,
description="Prompts separated by '|' (max 10)",
example="a cat|a dog|a bird"),
model: Literal["Turbo", "Raw"] = Query("Turbo"),
steps: int = Query(8, ge=1, le=50),
guidance: float = Query(0.0, ge=0.0, le=10.0),
resolution: Optional[Literal["square","portrait","landscape","square2k"]] = None,
async_mode: bool = Query(False, alias="async"),
workers: int = Query(DEFAULT_WORKERS, ge=10, le=200),
):
"""
**Batch generation** — different prompts in parallel.
Separate prompts with `|` character.
"""
prompt_list = [p.strip() for p in prompts.split("|") if p.strip()]
if not prompt_list:
raise HTTPException(400, "No valid prompts provided")
if len(prompt_list) > MAX_PARALLEL_JOBS:
raise HTTPException(400, f"Max {MAX_PARALLEL_JOBS} prompts per batch")
base_url = get_base_url(request)
try:
jobs = [build_job(
prompt=p, model=model, steps=steps, guidance=guidance,
resolution=resolution
) for p in prompt_list]
except ValueError as e:
raise HTTPException(400, str(e))
if async_mode:
for job in jobs:
background_tasks.add_task(execute_job, job, workers, base_url)
return {
"status": "queued",
"count": len(jobs),
"job_ids": [j.job_id for j in jobs],
"poll_urls": [f"{base_url}/jobs/{j.job_id}" for j in jobs],
}
with ThreadPoolExecutor(max_workers=len(jobs)) as pool:
futs = {pool.submit(execute_job, j, workers, base_url): j for j in jobs}
for f in as_completed(futs): pass
return {
"status": "batch_complete",
"total": len(jobs),
"success": sum(1 for j in jobs if j.status == "success"),
"results": [j.to_public_dict() for j in jobs],
}
# ─── JOBS ────────────────────────────────────────────────────
@app.get("/jobs")
def list_jobs(
status: Optional[Literal["pending","running","success","failed"]] = None,
limit: int = Query(50, ge=1, le=500),
):
"""List all jobs (optionally filtered by status)."""
with JOBS_LOCK:
items = list(JOBS.values())
items.sort(key=lambda j: j.created_at, reverse=True)
if status:
items = [j for j in items if j.status == status]
items = items[:limit]
return {
"total": len(JOBS),
"returned": len(items),
"jobs": [j.to_public_dict() for j in items],
}
@app.get("/jobs/{job_id}")
def get_job(job_id: str):
"""Get status/details of a single job by ID."""
with JOBS_LOCK:
job = JOBS.get(job_id)
if not job:
raise HTTPException(404, f"Job {job_id} not found")
return job.to_public_dict()
@app.get("/jobs/{job_id}/result")
def get_job_result(
job_id: str,
wait: bool = Query(True, description="Block until job completes (max 90s)"),
timeout: int = Query(90, ge=1, le=300),
):
"""
Get the result of a job. If `wait=true` and job is still running,
blocks until completion or timeout.
"""
with JOBS_LOCK:
job = JOBS.get(job_id)
if not job:
raise HTTPException(404, f"Job {job_id} not found")
if wait and job.status in ("pending", "running"):
deadline = time.time() + timeout
while time.time() < deadline and job.status in ("pending", "running"):
time.sleep(0.5)
if job.status in ("pending", "running"):
raise HTTPException(408, "Timeout waiting for job")
if job.status == "success":
return job.to_public_dict()
raise HTTPException(500, job.to_public_dict())
# ─── IMAGE DOWNLOAD ──────────────────────────────────────────
@app.get("/image/{filename}")
def get_image(filename: str):
"""Download a generated image PNG."""
# sanitize: solo nomi tipo job_xxx.png
if not re.match(r"^job_[a-zA-Z0-9_]+\.png$", filename):
raise HTTPException(400, "Invalid filename")
file_path = IMAGES_DIR / filename
if not file_path.exists():
raise HTTPException(404, "Image not found (may have expired)")
return FileResponse(file_path, media_type="image/png", filename=filename)
# ─── STATS & ADMIN ───────────────────────────────────────────
@app.get("/stats")
def stats():
"""Statistics: proxy DB + jobs registry."""
db_stats = DB.stats()
with JOBS_LOCK:
job_stats = {
"total": len(JOBS),
"pending": sum(1 for j in JOBS.values() if j.status == "pending"),
"running": sum(1 for j in JOBS.values() if j.status == "running"),
"success": sum(1 for j in JOBS.values() if j.status == "success"),
"failed": sum(1 for j in JOBS.values() if j.status == "failed"),
}
return {
"proxy_db": db_stats,
"jobs": job_stats,
"used_proxies_now": len(USED_PROXIES),
"uptime_s": time.time() - START_TIME,
}
@app.get("/admin/refresh")
def admin_refresh():
"""Force refresh of proxy pool from ProxyScrape."""
n = refresh_proxy_pool()
return {"status": "ok", "new_proxies_fetched": n, "db_stats": DB.stats()}
@app.get("/admin/cleanup")
def admin_cleanup():
"""Clean up old jobs from registry."""
n = cleanup_old_jobs()
return {"status": "ok", "jobs_removed": n, "jobs_remaining": len(JOBS)}
# ═══════════════════════════════════════════════════════════════
# OPENAPI TOOL SCHEMA (per agent AI)
# ═══════════════════════════════════════════════════════════════
@app.get("/tool-schema")
def tool_schema(request: Request):
"""
OpenAI/Anthropic-compatible tool schema for agent integration.
Copy this JSON directly into your agent's tool definitions.
"""
base = get_base_url(request)
return {
"openai_function_call": {
"type": "function",
"function": {
"name": "generate_image",
"description": (
"Generate an AI image from a text description using Krea-2 model. "
"Returns a URL where the image can be downloaded. "
"Fast (5-30 seconds), no login required."
),
"parameters": {
"type": "object",
"properties": {
"prompt": {"type": "string",
"description": "Detailed text description of the image to generate"},
"model": {"type": "string", "enum": ["Turbo", "Raw"],
"description": "Turbo=fast, Raw=high quality", "default": "Turbo"},
"steps": {"type": "integer", "minimum": 1, "maximum": 50,
"description": "Denoising steps", "default": 8},
"resolution": {"type": "string",
"enum": ["square", "portrait", "landscape", "square2k"],
"default": "square"},
"n": {"type": "integer", "minimum": 1, "maximum": 10,
"description": "Number of variants", "default": 1},
},
"required": ["prompt"],
},
},
"api_endpoint": f"GET {base}/generate",
},
"anthropic_tool": {
"name": "generate_image",
"description": "Generate AI images from text prompts via Krea-2",
"input_schema": {
"type": "object",
"properties": {
"prompt": {"type": "string"},
"model": {"type": "string", "enum": ["Turbo", "Raw"]},
"steps": {"type": "integer"},
"resolution": {"type": "string",
"enum": ["square", "portrait", "landscape", "square2k"]},
"n": {"type": "integer"},
},
"required": ["prompt"],
},
"api_endpoint": f"GET {base}/generate",
},
"langchain_tool_example": (
"from langchain.tools import tool\n"
"import requests\n\n"
"@tool\n"
"def generate_image(prompt: str, n: int = 1) -> str:\n"
" '''Generate AI images. Returns URLs.'''\n"
f" r = requests.get('{base}/generate', params={{'prompt': prompt, 'n': n}})\n"
" data = r.json()\n"
" if data['status'] == 'success':\n"
" return data['result']['local_url']\n"
" return [j['result']['local_url'] for j in data['results']]"
),
}
# ═══════════════════════════════════════════════════════════════
# STARTUP
# ═══════════════════════════════════════════════════════════════
START_TIME = time.time()
@app.on_event("startup")
async def startup():
"""Refresh iniziale del pool proxy."""
print(f"\n{'='*60}")
print(f" ⚡ KREA-2 API STARTING")
print(f"{'='*60}")
print(f" → Refreshing proxy pool from ProxyScrape...")
n = refresh_proxy_pool()
s = DB.stats()
print(f" → {n} new proxies fetched")
print(f" → DB: {s['total']} total ({s['good']} good, {s['banned']} banned)")
print(f" → Images directory: {IMAGES_DIR.absolute()}")
print(f" → Ready ✓\n")
# Cleanup periodico background
async def periodic_cleanup():
while True:
await asyncio.sleep(600) # ogni 10min
n = cleanup_old_jobs()
if n > 0:
print(f" [cleanup] removed {n} old jobs")
asyncio.create_task(periodic_cleanup())
# ═══════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════
def main():
parser = argparse.ArgumentParser(description="Krea-2 REST API server")
parser.add_argument("--host", default="0.0.0.0", help="Bind host")
parser.add_argument("--port", type=int, default=7860, help="Bind port")
parser.add_argument("--reload", action="store_true", help="Auto-reload on changes")
parser.add_argument("--workers", type=int, default=1,
help="Uvicorn workers (>1 disabilita stato in-memory condiviso)")
args = parser.parse_args()
if os.name == "nt": os.system("")
print(f"\n Starting on http://{args.host}:{args.port}")
print(f" Docs: http://{args.host}:{args.port}/docs")
print(f" OpenAPI: http://{args.host}:{args.port}/openapi.json\n")
uvicorn.run(
"api:app" if args.reload else app,
host=args.host,
port=args.port,
reload=args.reload,
workers=args.workers if not args.reload else 1,
log_level="info",
)
if __name__ == "__main__":
import os
if os.name == "nt": os.system("")
uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")