|
|
| """
|
| 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
|
|
|
|
|
|
|
|
|
|
|
| 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_TTL_SECONDS = 3600
|
|
|
| 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),
|
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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"
|
|
|
|
|
|
|
|
|
|
|
|
|
| @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"
|
| created_at: float = 0.0
|
| started_at: float = 0.0
|
| completed_at: float = 0.0
|
|
|
|
|
| filename: str = ""
|
| file_path: str = ""
|
| file_size: int = 0
|
| hf_url: str = ""
|
| local_url: str = ""
|
| result_seed: int = 0
|
| error: str = ""
|
|
|
|
|
| winner_proxy: str = ""
|
| winner_latency_ms: float = 0.0
|
| workers_launched: int = 0
|
| workers_dead: int = 0
|
| tokens_ok: int = 0
|
|
|
|
|
| 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
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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:
|
|
|
| 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"]}
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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()
|
|
|
|
|
| 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
|
|
|
|
|
| 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."""
|
|
|
| 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]
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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}"
|
|
|
|
|
|
|
|
|
| @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"""<!DOCTYPE html>
|
| <html><head><title>Krea-2 API</title>
|
| <style>
|
| body{{font-family:system-ui,-apple-system,sans-serif;max-width:900px;
|
| margin:2rem auto;padding:0 1rem;background:#0d0d1a;color:#e0e0f0;line-height:1.6}}
|
| h1{{background:linear-gradient(90deg,#00ffdc,#ff50dc);-webkit-background-clip:text;
|
| -webkit-text-fill-color:transparent;font-size:2.5rem;margin-bottom:0}}
|
| .subtitle{{color:#a0a0c0;margin-top:0}}
|
| code{{background:#1a1a2e;padding:.2rem .4rem;border-radius:3px;color:#00ffdc}}
|
| pre{{background:#1a1a2e;padding:1rem;border-radius:6px;overflow-x:auto;
|
| border-left:3px solid #ff50dc}}
|
| .endpoint{{background:#151525;padding:.8rem 1rem;margin:.5rem 0;
|
| border-radius:6px;border-left:3px solid #00ffdc}}
|
| .method{{color:#00ffdc;font-weight:bold}}
|
| a{{color:#ff50dc}}
|
| .stats{{display:grid;grid-template-columns:repeat(4,1fr);gap:1rem;margin:1rem 0}}
|
| .stat{{background:#151525;padding:1rem;border-radius:6px;text-align:center}}
|
| .stat b{{font-size:1.8rem;display:block;color:#00ffdc}}
|
| .badge{{display:inline-block;padding:.15rem .5rem;border-radius:3px;
|
| font-size:.75rem;background:#00ffdc;color:#0d0d1a;font-weight:bold}}
|
| </style></head><body>
|
| <h1>⚡ Krea-2 API</h1>
|
| <p class="subtitle">AI image generation as a service — anonymous, proxy-load-balanced, agent-ready</p>
|
|
|
| <div class="stats">
|
| <div class="stat"><b>{stats['total']}</b>proxies</div>
|
| <div class="stat"><b>{stats['good']}</b>good</div>
|
| <div class="stat"><b>{n_jobs}</b>total jobs</div>
|
| <div class="stat"><b>{n_running}</b>running</div>
|
| </div>
|
|
|
| <h2>🔌 Quick Start</h2>
|
|
|
| <div class="endpoint">
|
| <span class="method">GET</span> <code>/generate?prompt=a+cyberpunk+cat</code>
|
| <span class="badge">SYNC</span>
|
| <p>Blocca fino al completamento (~5-30s), ritorna JSON con URL immagine.</p>
|
| </div>
|
|
|
| <div class="endpoint">
|
| <span class="method">GET</span> <code>/generate?prompt=...&async=true</code>
|
| <span class="badge">ASYNC</span>
|
| <p>Ritorna subito <code>job_id</code>, poi polling con <code>/jobs/{{job_id}}</code>.</p>
|
| </div>
|
|
|
| <div class="endpoint">
|
| <span class="method">GET</span> <code>/generate?prompt=...&n=5</code>
|
| <span class="badge">BATCH</span>
|
| <p>Genera N varianti in parallelo (max 10).</p>
|
| </div>
|
|
|
| <h2>📋 All Endpoints</h2>
|
| <pre>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</pre>
|
|
|
| <h2>🎯 Example — Agent tool call</h2>
|
| <pre>curl "{base}/generate?prompt=sunset&model=Turbo&steps=8"</pre>
|
|
|
| <pre>{{
|
| "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
|
| }}
|
| }}</pre>
|
|
|
| <p style="margin-top:2rem;color:#666;text-align:center">
|
| <a href="/docs">Swagger docs</a> · <a href="/openapi.json">OpenAPI</a> · <a href="/stats">Stats</a>
|
| </p>
|
| </body></html>
|
| """
|
|
|
|
|
| @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,
|
| }
|
|
|
|
|
|
|
|
|
| @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:
|
|
|
| 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",
|
| })
|
|
|
|
|
| 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())
|
|
|
|
|
| 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],
|
| }
|
|
|
|
|
|
|
|
|
| @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())
|
|
|
|
|
|
|
|
|
| @app.get("/image/{filename}")
|
| def get_image(filename: str):
|
| """Download a generated image 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)
|
|
|
|
|
|
|
|
|
| @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)}
|
|
|
|
|
|
|
|
|
|
|
|
|
| @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']]"
|
| ),
|
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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")
|
|
|
|
|
| async def periodic_cleanup():
|
| while True:
|
| await asyncio.sleep(600)
|
| n = cleanup_old_jobs()
|
| if n > 0:
|
| print(f" [cleanup] removed {n} old jobs")
|
|
|
| asyncio.create_task(periodic_cleanup())
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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") |
