573 lines
30 KiB
Python
573 lines
30 KiB
Python
"""Private adapters for Huobao: Dify, Qwen Image, local H3, and Grok CLI."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
from contextlib import closing
|
|
from email import policy
|
|
from email.parser import BytesParser
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import secrets
|
|
import sqlite3
|
|
import struct
|
|
import tempfile
|
|
import time
|
|
from urllib.parse import quote
|
|
|
|
import httpx
|
|
from fastapi import Depends, FastAPI, Header, HTTPException, Request
|
|
from fastapi.responses import JSONResponse, StreamingResponse
|
|
|
|
|
|
def chat_query(body):
|
|
messages = body.get("messages")
|
|
if not isinstance(messages, list) or not messages:
|
|
raise HTTPException(400, "messages must be a non-empty array")
|
|
if body.get("tool_choice") == "required" or isinstance(body.get("tool_choice"), dict):
|
|
raise HTTPException(400, "This chatflow does not support required tool calls; use DeepSeek")
|
|
parts = []
|
|
for item in messages:
|
|
content = item.get("content", "")
|
|
if isinstance(content, list):
|
|
if any(p.get("type") != "text" for p in content):
|
|
raise HTTPException(400, "This adapter accepts text messages only")
|
|
content = "\n".join(p.get("text", "") for p in content)
|
|
if content:
|
|
parts.append(f"[{item.get('role', 'user')}]\n{content}")
|
|
fmt = body.get("response_format") or {}
|
|
if fmt.get("type") == "json_schema":
|
|
parts.append("Return only JSON matching this schema:\n" + json.dumps(fmt.get("json_schema", {}).get("schema", {}), ensure_ascii=False))
|
|
elif fmt.get("type") == "json_object":
|
|
parts.append("Return only a valid JSON object, without Markdown fences.")
|
|
if not parts:
|
|
raise HTTPException(400, "Empty messages")
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
def image_dimensions(data_url):
|
|
try:
|
|
data = base64.b64decode(data_url.split(",", 1)[1], validate=True)
|
|
if data.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
return struct.unpack(">II", data[16:24])
|
|
if data.startswith((b"GIF87a", b"GIF89a")):
|
|
return struct.unpack("<HH", data[6:10])
|
|
if data[:2] == b"\xff\xd8":
|
|
offset = 2
|
|
while offset + 4 < len(data):
|
|
if data[offset] != 255:
|
|
offset += 1
|
|
continue
|
|
while data[offset] == 255:
|
|
offset += 1
|
|
marker = data[offset]
|
|
offset += 1
|
|
if marker in {0xD8, 0xD9} or 0xD0 <= marker <= 0xD7:
|
|
continue
|
|
size = int.from_bytes(data[offset:offset + 2], "big")
|
|
if marker in {0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF}:
|
|
height, width = struct.unpack(">HH", data[offset + 3:offset + 7])
|
|
return width, height
|
|
if size < 2:
|
|
break
|
|
offset += size
|
|
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
|
if data[12:16] == b"VP8X":
|
|
return 1 + int.from_bytes(data[24:27], "little"), 1 + int.from_bytes(data[27:30], "little")
|
|
if data[12:16] == b"VP8L" and data[20] == 0x2F:
|
|
bits = int.from_bytes(data[21:25], "little")
|
|
return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
|
|
if data[12:16] == b"VP8 " and data[23:26] == b"\x9d\x01\x2a":
|
|
width, height = struct.unpack("<HH", data[26:30])
|
|
return width & 0x3FFF, height & 0x3FFF
|
|
except (ValueError, IndexError, struct.error):
|
|
pass
|
|
raise HTTPException(400, "Unable to read first-frame dimensions; use PNG, JPEG or WebP")
|
|
|
|
|
|
VIDEO_MODELS = {
|
|
"MiniMax-H3-VDN-H3": {"reference": "minimax_h3_vdn_hd", "min_images": 1, "max_images": 9},
|
|
"MiniMax-H3": {"text": "minimax_h3_t2v_official", "first": "minimax_h3_i2v", "reference": "minimax_h3_r2v", "min_images": 0, "max_images": 9},
|
|
"MiniMax-H3-Hybrid": {"text": "minimax_h3_hybrid_v0331_t2v", "first": "minimax_h3_hybrid_v0331_i2v", "reference": "minimax_h3_hybrid_v0331_i2v", "min_images": 0, "max_images": 9},
|
|
"MiniMax-H3-R2V-4Step": {"reference": "minimax_h3_r2v_4step", "min_images": 1, "max_images": 9},
|
|
"LTX-2.3": {"text": "ltx23_t2v", "first": "ltx23_t2v", "reference": "ltx23_t2v", "min_images": 0, "max_images": 1},
|
|
"LTX-2.3-Text-Preserve": {"text": "ltx23_t2v_text_preserve", "first": "ltx23_t2v_text_preserve", "reference": "ltx23_t2v_text_preserve", "min_images": 0, "max_images": 1},
|
|
"Wan2.2-I2V": {"first": "wan22_simple", "reference": "wan22_simple", "min_images": 1, "max_images": 1},
|
|
}
|
|
|
|
|
|
def video_workflows(spec):
|
|
return {spec[mode] for mode in ("text", "first", "reference") if mode in spec}
|
|
|
|
|
|
def gpu_routing(body):
|
|
gpu = body.get("gpu", "auto")
|
|
if gpu not in ("auto", "gpu1", "gpu2", "gpu3"):
|
|
raise HTTPException(400, "gpu must be auto, gpu1, gpu2 or gpu3")
|
|
if gpu == "auto":
|
|
return {"provider_id": "default"}
|
|
return {"provider_id": "default" if gpu == "gpu2" else gpu, "strict_provider": True}
|
|
|
|
|
|
def h3_request(body):
|
|
"""Translate the local video protocol. 'default' means upstream auto-balance."""
|
|
model = body.get("model", "MiniMax-H3")
|
|
spec = VIDEO_MODELS.get(model)
|
|
if spec is None:
|
|
raise HTTPException(400, "Unknown local video model; query /v1/models")
|
|
content = body.get("content") or []
|
|
if not isinstance(content, list):
|
|
raise HTTPException(400, "content must be an array")
|
|
prompt = "\n".join(str(c.get("text", "")) for c in content if c.get("type") == "text").strip()
|
|
if not prompt or len(prompt) > 7000:
|
|
raise HTTPException(400, "A prompt of 1-7000 characters is required")
|
|
images, roles = [], []
|
|
for item in content:
|
|
kind = item.get("type")
|
|
if kind == "text":
|
|
continue
|
|
if kind != "image_url" or item.get("role") == "last_frame":
|
|
raise HTTPException(400, "Local video supports first-frame/reference images; last frames, reference video and voice samples are unsupported")
|
|
raw = item.get("image_url") or {}
|
|
url = raw.get("url", "") if isinstance(raw, dict) else raw
|
|
if not isinstance(url, str) or not url.startswith("data:image/"):
|
|
raise HTTPException(400, "Local video requires inline image data; enable Huobao's local-image conversion")
|
|
images.append(url)
|
|
roles.append(item.get("role", "reference_image"))
|
|
if not spec["min_images"] <= len(images) <= spec["max_images"]:
|
|
raise HTTPException(400, f"{model} requires {spec['min_images']}-{spec['max_images']} images")
|
|
if "first_frame" in roles and len(images) != 1:
|
|
raise HTTPException(400, "Use either one fixed first frame or reference images; these modes cannot be combined")
|
|
mode = "text" if not images else "first" if roles == ["first_frame"] else "reference"
|
|
if mode not in spec:
|
|
raise HTTPException(400, f"{model} does not support {mode} mode; use reference images")
|
|
try:
|
|
duration = float(body.get("duration", 5))
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(400, "Invalid duration")
|
|
if not math.isfinite(duration) or not 4 <= duration <= 15:
|
|
raise HTTPException(400, "Duration must be 4-15 seconds")
|
|
ratio = body.get("ratio", "16:9")
|
|
if ratio not in {"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}:
|
|
raise HTTPException(400, "Unsupported aspect ratio")
|
|
resolution = str(body.get("resolution", "768P")).upper()
|
|
resolutions = {"480P": 480, "720P": 720, "1080P": 1080, "768P": 768, "2K": 1152}
|
|
if model == "MiniMax-H3-VDN-H3":
|
|
if ratio not in {"16:9", "9:16"}:
|
|
raise HTTPException(400, "VDN HD supports 16:9 or 9:16")
|
|
resolutions = {"2K": 1440, "4K": 2160}
|
|
if resolution not in resolutions:
|
|
raise HTTPException(400, "Unsupported video resolution")
|
|
short = resolutions[resolution]
|
|
a, b = map(int, ratio.split(":"))
|
|
if roles == ["first_frame"] and "ratio" not in body:
|
|
a, b = image_dimensions(images[0])
|
|
if a <= 0 or b <= 0:
|
|
raise HTTPException(400, "Invalid first-frame dimensions")
|
|
width = round(short * max(1, a / b) / 64) * 64
|
|
height = round(short * max(1, b / a) / 64) * 64
|
|
if model == "MiniMax-H3-VDN-H3":
|
|
width, height = (short * 16 // 9, short) if a > b else (short, short * 16 // 9)
|
|
workflow = spec[mode]
|
|
request = {**gpu_routing(body), "workflow_key": workflow, "prompt": prompt,
|
|
"input_images": images, "duration": duration, "width": width, "height": height,
|
|
"metadata": {"source": "huobao-drama", "model": model}}
|
|
if body.get("seed") is not None:
|
|
if not isinstance(body["seed"], int) or not 0 <= body["seed"] < 2**63:
|
|
raise HTTPException(400, "Invalid seed")
|
|
request["seed"] = body["seed"]
|
|
return request
|
|
|
|
|
|
def qwen_image_request(body, images=None):
|
|
if body.get("model") != "qwen_image_2_1":
|
|
raise HTTPException(400, "Use model qwen_image_2_1")
|
|
prompt = body.get("prompt")
|
|
if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 32000:
|
|
raise HTTPException(400, "A prompt of 1-32000 characters is required")
|
|
if str(body.get("n", 1)) != "1":
|
|
raise HTTPException(400, "This image adapter supports n=1 per task")
|
|
if body.get("response_format", "url") != "url":
|
|
raise HTTPException(400, "Asynchronous image tasks require response_format=url")
|
|
size = str(body.get("size", "1024x1024")).lower()
|
|
if size == "auto":
|
|
size = "1024x1024"
|
|
match = re.fullmatch(r"(\d{3,4})x(\d{3,4})", size)
|
|
if not match or not all(256 <= int(d) <= 4096 for d in match.groups()):
|
|
raise HTTPException(400, "size must be WIDTHxHEIGHT with dimensions 256-4096")
|
|
width, height = map(int, match.groups())
|
|
refs = images or []
|
|
if len(refs) > 6:
|
|
raise HTTPException(400, "Qwen Image 2.1 supports at most 6 reference images")
|
|
return {**gpu_routing(body), "workflow_key": "qwen_image_2_1", "prompt": prompt.strip(),
|
|
"width": width, "height": height, "batch_size": 1, "input_images": refs,
|
|
"metadata": {"source": "huobao-drama", "output_type": "image"}}
|
|
|
|
|
|
async def parse_image_edit(request):
|
|
content_type = request.headers.get("content-type", "")
|
|
if not content_type.lower().startswith("multipart/form-data;"):
|
|
raise HTTPException(415, "Image edits require multipart/form-data")
|
|
raw = bytearray()
|
|
async for chunk in request.stream():
|
|
raw.extend(chunk)
|
|
if len(raw) > 64 * 1024 * 1024:
|
|
raise HTTPException(413, "Reference uploads exceed 64 MiB")
|
|
message = BytesParser(policy=policy.default).parsebytes(
|
|
("Content-Type: " + content_type + "\r\nMIME-Version: 1.0\r\n\r\n").encode() + raw)
|
|
if not message.is_multipart() or message.defects:
|
|
raise HTTPException(400, "Invalid multipart image request")
|
|
fields, images = {}, []
|
|
for part in message.iter_parts():
|
|
name = part.get_param("name", header="content-disposition")
|
|
data = part.get_payload(decode=True)
|
|
if not isinstance(data, bytes):
|
|
raise HTTPException(400, "Invalid multipart field")
|
|
if name in {"image", "image[]"}:
|
|
mime = part.get_content_type()
|
|
if mime not in {"image/png", "image/jpeg", "image/webp"} or not data or len(data) > 10 * 1024 * 1024:
|
|
raise HTTPException(400, "Each reference must be PNG, JPEG or WebP, up to 10 MiB")
|
|
url = "data:" + mime + ";base64," + base64.b64encode(data).decode()
|
|
width, height = image_dimensions(url)
|
|
if not (0 < width <= 16384 and 0 < height <= 16384):
|
|
raise HTTPException(400, "Invalid reference image dimensions")
|
|
images.append(url)
|
|
if len(images) > 6:
|
|
raise HTTPException(400, "Qwen Image 2.1 supports at most 6 reference images")
|
|
elif name in {"model", "prompt", "size", "n", "response_format", "gpu"}:
|
|
if name in fields or len(data) > 128000:
|
|
raise HTTPException(400, "Duplicate or oversized form field")
|
|
try:
|
|
fields[name] = data.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
raise HTTPException(400, "Form text must be UTF-8")
|
|
else:
|
|
raise HTTPException(400, "Unsupported image edit field: " + str(name)[:60])
|
|
if not images:
|
|
raise HTTPException(400, "At least one reference image is required")
|
|
return qwen_image_request(fields, images)
|
|
|
|
|
|
def create_app(config):
|
|
app = FastAPI(title="Huobao private model bridge", docs_url=None, redoc_url=None)
|
|
state_dir = Path(config.get("state_dir", "./state"))
|
|
state_dir.mkdir(parents=True, exist_ok=True)
|
|
db_path = state_dir / "jobs.sqlite3"
|
|
with closing(sqlite3.connect(db_path)) as db, db:
|
|
db.execute("CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, duration REAL NOT NULL, created REAL NOT NULL)")
|
|
db.execute("CREATE TABLE IF NOT EXISTS image_jobs (id TEXT PRIMARY KEY, created REAL NOT NULL)")
|
|
|
|
def auth(authorization: str | None = Header(default=None)):
|
|
expected = config.get("token", "")
|
|
if not expected or not secrets.compare_digest(authorization or "", "Bearer " + expected):
|
|
raise HTTPException(401, "Invalid bridge token")
|
|
|
|
@app.exception_handler(HTTPException)
|
|
async def http_error(request, exc):
|
|
return JSONResponse({"error": {"message": str(exc.detail), "type": "bridge_error", "code": exc.status_code}}, status_code=exc.status_code)
|
|
|
|
@app.exception_handler(httpx.HTTPError)
|
|
async def upstream_error(request, exc):
|
|
return JSONResponse({"error": {"message": "Upstream connection failed: " + type(exc).__name__, "type": "upstream_error"}}, status_code=502)
|
|
|
|
def upstream_headers():
|
|
return {"Authorization": "Bearer " + config["h3"]["token"]}
|
|
|
|
async def upstream(method, path, body=None):
|
|
if not config.get("h3"):
|
|
raise HTTPException(404, "H3 is not configured on this bridge")
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(180, connect=10), trust_env=False) as client:
|
|
r = await client.request(method, config["h3"]["url"] + path, headers=upstream_headers(), json=body)
|
|
if r.is_error:
|
|
try:
|
|
detail = r.json().get("detail", "Model service request failed")
|
|
except ValueError:
|
|
detail = "Model service request failed"
|
|
raise HTTPException(r.status_code, str(detail)[:600])
|
|
return r.json()
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "ok", "service": "huobao-model-bridge"}
|
|
|
|
@app.get("/v1/models", dependencies=[Depends(auth)])
|
|
async def models():
|
|
names = list(config.get("dify", {}))
|
|
if config.get("h3"):
|
|
# Verify the real service, instead of making a local catalogue look healthy.
|
|
registry = await upstream("GET", "/v1/models")
|
|
available = {m["id"] for m in registry.get("data", [])}
|
|
names.extend(name for name, spec in VIDEO_MODELS.items() if video_workflows(spec) <= available)
|
|
if any(m.get("id") == "qwen_image_2_1" and m.get("model_type") == "image" for m in registry.get("data", [])):
|
|
names.append("qwen_image_2_1")
|
|
if config.get("grok"):
|
|
names.append("grok-local")
|
|
return {"object": "list", "data": [{"id": n, "object": "model", "owned_by": "private"} for n in names]}
|
|
|
|
async def dify_events(model, query):
|
|
upstream_config = config["dify"][model]
|
|
payload = {"inputs": {}, "query": query, "response_mode": "streaming",
|
|
"user": "huobao-" + secrets.token_hex(8), "auto_generate_name": False}
|
|
finished, usage = False, {}
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(240, connect=10), trust_env=False) as client:
|
|
async with client.stream("POST", upstream_config["url"] + "/chat-messages",
|
|
headers={"Authorization": "Bearer " + upstream_config["key"]}, json=payload) as response:
|
|
if response.is_error:
|
|
await response.aread()
|
|
try:
|
|
message = response.json().get("message", "Dify request failed")
|
|
except ValueError:
|
|
message = "Dify request failed"
|
|
raise HTTPException(response.status_code, str(message)[:400])
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith("data:"):
|
|
continue
|
|
event = json.loads(line[5:].strip())
|
|
kind = event.get("event")
|
|
if kind in {"message", "agent_message"}:
|
|
yield "text", event.get("answer", "")
|
|
elif kind == "error":
|
|
raise HTTPException(502, str(event.get("message", "Dify generation failed"))[:400])
|
|
elif kind == "message_end":
|
|
finished = True
|
|
usage = event.get("metadata", {}).get("usage", {})
|
|
elif kind == "workflow_finished":
|
|
data = event.get("data") or {}
|
|
if data.get("status") not in {None, "succeeded"}:
|
|
raise HTTPException(502, str(data.get("error") or "Dify workflow failed")[:400])
|
|
if not finished:
|
|
raise HTTPException(502, "Dify stream ended before message_end")
|
|
yield "usage", usage
|
|
|
|
async def grok_events(query):
|
|
item = config["grok"]
|
|
args = [item["exe"], "--cwd", item["cwd"], "--output-format", "json", "--verbatim",
|
|
"--no-subagents", "--no-memory", "--disable-web-search", "--max-turns", "1",
|
|
"--tools", "", "--deny", "MCPTool", "--permission-mode", "dontAsk",
|
|
"--system-prompt-override", "Answer the supplied writing request using text only. Do not use tools."]
|
|
options = {"creationflags": 0x08000000} if os.name == "nt" else {}
|
|
# A file avoids the Windows command-line length limit for full scripts.
|
|
with tempfile.TemporaryDirectory(prefix="huobao-grok-") as scratch:
|
|
prompt_file = Path(scratch) / "prompt.txt"
|
|
prompt_file.write_text(query, encoding="utf-8")
|
|
args += ["--prompt-file", str(prompt_file)]
|
|
process = await asyncio.create_subprocess_exec(*args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, **options)
|
|
try:
|
|
out, err = await asyncio.wait_for(process.communicate(), timeout=240)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
process.kill()
|
|
await process.wait()
|
|
raise
|
|
try:
|
|
data = json.loads(out.decode("utf-8"))
|
|
except (ValueError, UnicodeDecodeError):
|
|
raise HTTPException(502, "Grok CLI did not return JSON; run Grok login on this computer")
|
|
if process.returncode or not data.get("text"):
|
|
raise HTTPException(502, "Grok CLI failed; check local login and available quota")
|
|
yield "text", data["text"]
|
|
yield "usage", data.get("usage", {})
|
|
|
|
@app.post("/v1/chat/completions", dependencies=[Depends(auth)])
|
|
async def chat(body: dict):
|
|
model = body.get("model", "")
|
|
if model not in config.get("dify", {}) and not (model == "grok-local" and config.get("grok")):
|
|
raise HTTPException(404, "Unknown text model")
|
|
query = chat_query(body)
|
|
events = grok_events(query) if model == "grok-local" else dify_events(model, query)
|
|
cid, created = "chatcmpl-" + secrets.token_hex(12), int(time.time())
|
|
|
|
def chunk(delta, finish=None):
|
|
return {"id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
|
|
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
|
|
|
|
if body.get("stream"):
|
|
async def stream():
|
|
def encode(value):
|
|
return "data: " + json.dumps(value, ensure_ascii=False) + "\n\n"
|
|
yield encode(chunk({"role": "assistant", "content": ""}))
|
|
try:
|
|
async for kind, value in events:
|
|
if kind == "text":
|
|
yield encode(chunk({"content": value}))
|
|
yield encode(chunk({}, "stop"))
|
|
except (HTTPException, httpx.HTTPError) as error:
|
|
message = str(error.detail) if isinstance(error, HTTPException) else "Upstream connection failed"
|
|
yield encode({"error": {"message": message, "type": "upstream_error"}})
|
|
yield "data: [DONE]\n\n"
|
|
return StreamingResponse(stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache"})
|
|
text, usage = "", {}
|
|
async for kind, value in events:
|
|
if kind == "text":
|
|
text += value
|
|
else:
|
|
usage = value
|
|
return {"id": cid, "object": "chat.completion", "created": created, "model": model,
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
|
|
"usage": {"prompt_tokens": usage.get("prompt_tokens", usage.get("input_tokens", 0)),
|
|
"completion_tokens": usage.get("completion_tokens", usage.get("output_tokens", 0)),
|
|
"total_tokens": usage.get("total_tokens", 0)}}
|
|
|
|
async def require_qwen_image():
|
|
registry = await upstream("GET", "/v1/models")
|
|
if not any(m.get("id") == "qwen_image_2_1" and m.get("model_type") == "image" for m in registry.get("data", [])):
|
|
raise HTTPException(503, "Qwen Image 2.1 workflow is unavailable")
|
|
|
|
@app.get("/v1/images/pool", dependencies=[Depends(auth)])
|
|
async def image_pool():
|
|
await require_qwen_image()
|
|
health, catalogue = await asyncio.gather(upstream("GET", "/health"), upstream("GET", "/v1/providers"))
|
|
return {"model": "qwen_image_2_1", "auto_balance": health.get("autoBalance") is True,
|
|
"routing": "shortest-queue-with-reservations" if health.get("autoBalance") is True else "single-provider",
|
|
"configured_workers": [{"id": p["id"], "name": p.get("name", p["id"])} for p in catalogue.get("providers", [])]}
|
|
|
|
async def submit_image(payload):
|
|
await require_qwen_image()
|
|
result = await upstream("POST", "/v1/jobs", payload)
|
|
jid = result.get("id")
|
|
if not isinstance(jid, str) or not jid:
|
|
raise HTTPException(502, "Image service did not return a task ID")
|
|
with closing(sqlite3.connect(db_path)) as db, db:
|
|
db.execute("INSERT INTO image_jobs VALUES(?,?)", (jid, time.time()))
|
|
return {"task_id": jid, "status": "processing", "provider_id": result.get("providerId")}
|
|
|
|
@app.post("/v1/images/validate", dependencies=[Depends(auth)])
|
|
async def validate_image(body: dict):
|
|
payload = qwen_image_request(body)
|
|
await require_qwen_image()
|
|
return {"valid": True, "workflow": payload["workflow_key"], "requested_size": [payload["width"], payload["height"]], "max_reference_images": 6,
|
|
"pool": await image_pool()}
|
|
|
|
@app.post("/v1/images/generations", dependencies=[Depends(auth)])
|
|
async def generate_image(body: dict):
|
|
return await submit_image(qwen_image_request(body))
|
|
|
|
@app.post("/v1/images/edits", dependencies=[Depends(auth)])
|
|
async def edit_image(request: Request):
|
|
return await submit_image(await parse_image_edit(request))
|
|
|
|
@app.get("/v1/images/task/{jid}", dependencies=[Depends(auth)])
|
|
async def poll_image(jid: str):
|
|
with closing(sqlite3.connect(db_path)) as db:
|
|
found = db.execute("SELECT 1 FROM image_jobs WHERE id=?", (jid,)).fetchone()
|
|
if not found:
|
|
raise HTTPException(404, "Unknown image task")
|
|
result = await upstream("GET", "/v1/jobs/" + quote(jid, safe=""))
|
|
state = result.get("status", "running")
|
|
if state in {"failed", "cancelled"}:
|
|
return {"status": "failed", "error": {"message": str(result.get("error") or "Image generation failed")[:600]}}
|
|
if state != "succeeded":
|
|
return {"status": "processing", "provider_id": result.get("providerId")}
|
|
output = next((o for o in result.get("outputs", []) if o.get("kind") == "image"), None)
|
|
if output is None:
|
|
return {"status": "failed", "error": {"message": "Task completed without an image output"}}
|
|
index, expiry = int(output["index"]), int(time.time()) + 86400
|
|
extension = {"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp"}.get(output.get("mimeType"), "png")
|
|
# Huobao determines the saved extension from the URL, not Content-Type.
|
|
url = f"{config['public_url']}/media/{quote(jid, safe='')}/{index}/image.{extension}?expires={expiry}&sig={signature(jid, index, expiry)}"
|
|
return {"status": "completed", "image_url": url, "data": [{"url": url}], "provider_id": result.get("providerId")}
|
|
|
|
@app.post("/v2/video_generation/validate", dependencies=[Depends(auth)])
|
|
async def validate_video(body: dict):
|
|
request = h3_request(body)
|
|
registry = await upstream("GET", "/v1/models")
|
|
if request["workflow_key"] not in {m["id"] for m in registry.get("data", [])}:
|
|
raise HTTPException(503, "Selected video workflow is unavailable")
|
|
plan = await upstream("POST", "/v1/audio-plan", request)
|
|
return {"valid": plan.get("status", "ready") == "ready", "workflow": request["workflow_key"],
|
|
"width": request["width"], "height": request["height"], "duration": request["duration"], "audio_plan": plan,
|
|
"routing": "upstream-auto-balance"}
|
|
|
|
@app.post("/v2/video_generation", dependencies=[Depends(auth)])
|
|
async def video(body: dict):
|
|
# Huobao probes video providers with an empty body. Check the real
|
|
# upstream without submitting a GPU task, then report readiness.
|
|
if not body:
|
|
registry = await upstream("GET", "/v1/models")
|
|
available = {m["id"] for m in registry.get("data", [])}
|
|
if not {"minimax_h3_t2v_official", "minimax_h3_i2v", "minimax_h3_r2v"} <= available:
|
|
raise HTTPException(503, "Required H3 workflows are unavailable")
|
|
return {"status": "ready", "probe": True, "model": "MiniMax-H3"}
|
|
request = h3_request(body)
|
|
registry = await upstream("GET", "/v1/models")
|
|
if request["workflow_key"] not in {m["id"] for m in registry.get("data", [])}:
|
|
raise HTTPException(503, "Selected video workflow is unavailable")
|
|
result = await upstream("POST", "/v1/jobs", request)
|
|
jid = result.get("id")
|
|
if not isinstance(jid, str) or not jid:
|
|
raise HTTPException(502, "Video service did not return a job ID")
|
|
with closing(sqlite3.connect(db_path)) as db, db:
|
|
db.execute("INSERT INTO jobs VALUES(?,?,?)", (jid, request["duration"], time.time()))
|
|
return {"task_id": jid, "status": "submitted", "provider_id": result.get("providerId")}
|
|
|
|
def get_job(jid):
|
|
with closing(sqlite3.connect(db_path)) as db:
|
|
row = db.execute("SELECT duration FROM jobs WHERE id=?", (jid,)).fetchone()
|
|
if not row:
|
|
raise HTTPException(404, "Unknown bridge job")
|
|
return row[0]
|
|
|
|
def signature(jid, index, expiry):
|
|
data = f"{jid}:{index}:{expiry}".encode()
|
|
return hmac.new(config["token"].encode(), data, hashlib.sha256).hexdigest()
|
|
|
|
@app.get("/v2/query/video_generation/{jid}", dependencies=[Depends(auth)])
|
|
async def poll(jid: str):
|
|
duration = get_job(jid)
|
|
result = await upstream("GET", "/v1/jobs/" + quote(jid, safe=""))
|
|
state = result.get("status", "running")
|
|
task = {"id": jid, "status": state, "duration": duration, "provider_id": result.get("providerId")}
|
|
if state == "succeeded":
|
|
output = next((o for o in result.get("outputs", []) if o.get("kind") == "video"), None)
|
|
if not output:
|
|
task.update(status="failed", error="Video task completed without a video output")
|
|
else:
|
|
index, expiry = int(output["index"]), int(time.time()) + 86400
|
|
url = f"{config['public_url']}/media/{quote(jid, safe='')}/{index}/video.mp4?expires={expiry}&sig={signature(jid, index, expiry)}"
|
|
task["content"] = {"url": url}
|
|
elif state in {"failed", "cancelled"}:
|
|
task["error"] = str(result.get("error") or "Video task failed")[:600]
|
|
return {"task": task}
|
|
|
|
@app.get("/media/{jid}/{index}")
|
|
@app.get("/media/{jid}/{index}/{filename}")
|
|
async def media(jid: str, index: int, expires: int, sig: str, filename: str | None = None):
|
|
if filename is not None and not re.fullmatch(r"image\.(png|jpg|webp)|video\.mp4", filename):
|
|
raise HTTPException(404, "Unknown media filename")
|
|
with closing(sqlite3.connect(db_path)) as db:
|
|
found = db.execute("SELECT 1 FROM jobs WHERE id=? UNION ALL SELECT 1 FROM image_jobs WHERE id=?", (jid, jid)).fetchone()
|
|
if not found:
|
|
raise HTTPException(404, "Unknown bridge task")
|
|
if expires < time.time() or not hmac.compare_digest(sig, signature(jid, index, expires)):
|
|
raise HTTPException(403, "Invalid or expired media link; poll the task again")
|
|
client = httpx.AsyncClient(timeout=httpx.Timeout(180, connect=10), trust_env=False)
|
|
url = config["h3"]["url"] + f"/v1/jobs/{quote(jid, safe='')}/outputs/{index}"
|
|
try:
|
|
response = await client.send(client.build_request("GET", url, headers=upstream_headers()), stream=True)
|
|
if response.is_error:
|
|
await response.aclose()
|
|
await client.aclose()
|
|
raise HTTPException(response.status_code, "Model output is not available")
|
|
except BaseException:
|
|
await client.aclose()
|
|
raise
|
|
async def download():
|
|
try:
|
|
async for part in response.aiter_bytes():
|
|
yield part
|
|
finally:
|
|
await response.aclose()
|
|
await client.aclose()
|
|
return StreamingResponse(download(), media_type=response.headers.get("content-type", "video/mp4"))
|
|
|
|
return app
|
|
|
|
|
|
def create_from_file():
|
|
path = Path(os.environ.get("HUOBAO_BRIDGE_CONFIG", "/bridge/config.json"))
|
|
return create_app(json.loads(path.read_text(encoding="utf-8")))
|