Initial SQL-only 1C adapter baseline
This commit is contained in:
@@ -0,0 +1,930 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import base64
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
MODEL = None
|
||||
TOKENIZER = None
|
||||
PROCESSOR = None
|
||||
EDIT_MODEL = None
|
||||
SERVER_CONFIG: dict = {}
|
||||
LOAD_ERROR: str | None = None
|
||||
IMAGE_JOBS: dict[str, dict] = {}
|
||||
IMAGE_JOB_LOCK = threading.Lock()
|
||||
IMAGE_JOB_CANCEL_EVENTS: dict[str, threading.Event] = {}
|
||||
IMAGE_MODEL_LOCK = threading.Lock()
|
||||
IMAGE_EDIT_MODEL_LOCK = threading.Lock()
|
||||
MODEL_LOAD_STATE: dict[str, dict] = {
|
||||
"image": {"status": "not_loaded"},
|
||||
"image_edit": {"status": "not_loaded"},
|
||||
}
|
||||
|
||||
|
||||
def set_load_state(key: str, **fields: object) -> None:
|
||||
state = MODEL_LOAD_STATE.setdefault(key, {})
|
||||
state.update(fields)
|
||||
state["updated_at"] = time.time()
|
||||
|
||||
|
||||
def release_cuda_memory() -> None:
|
||||
gc.collect()
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def unload_image_model() -> None:
|
||||
global MODEL
|
||||
if MODEL is None:
|
||||
return
|
||||
MODEL = None
|
||||
release_cuda_memory()
|
||||
set_load_state("image", status="not_loaded", unloaded_at=time.time())
|
||||
|
||||
|
||||
def unload_image_edit_model() -> None:
|
||||
global EDIT_MODEL
|
||||
if EDIT_MODEL is None:
|
||||
return
|
||||
EDIT_MODEL = None
|
||||
release_cuda_memory()
|
||||
set_load_state("image_edit", status="not_loaded", unloaded_at=time.time())
|
||||
|
||||
|
||||
def has_active_image_job(operation: str) -> bool:
|
||||
active_statuses = {"queued", "loading_model", "running", "cancel_requested"}
|
||||
with IMAGE_JOB_LOCK:
|
||||
return any(
|
||||
job.get("operation") == operation and job.get("status") in active_statuses
|
||||
for job in IMAGE_JOBS.values()
|
||||
)
|
||||
|
||||
|
||||
def diffusers_class_name(model_path: str | None) -> str:
|
||||
if not model_path:
|
||||
return ""
|
||||
index_path = Path(model_path) / "model_index.json"
|
||||
if not index_path.exists():
|
||||
return ""
|
||||
try:
|
||||
data = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return ""
|
||||
return str(data.get("_class_name") or "")
|
||||
|
||||
|
||||
def is_qwen_image_class(class_name: str) -> bool:
|
||||
return class_name in {"QwenImagePipeline", "QwenImageEditPipeline"}
|
||||
|
||||
|
||||
def has_fp16_variant(model_path: str) -> bool:
|
||||
path = Path(model_path)
|
||||
if not path.exists():
|
||||
return False
|
||||
return any(path.rglob("*fp16.safetensors"))
|
||||
|
||||
|
||||
def prepare_diffusers_pipeline(pipe, *, torch_module, class_name: str) -> None:
|
||||
if hasattr(pipe, "enable_attention_slicing"):
|
||||
pipe.enable_attention_slicing()
|
||||
if not torch_module.cuda.is_available():
|
||||
return
|
||||
if is_qwen_image_class(class_name) and hasattr(pipe, "enable_model_cpu_offload"):
|
||||
pipe.enable_model_cpu_offload()
|
||||
else:
|
||||
pipe.to("cuda")
|
||||
|
||||
|
||||
def load_diffusers_pipeline(model_path: str, *, edit: bool):
|
||||
import torch
|
||||
|
||||
class_name = diffusers_class_name(model_path)
|
||||
dtype = torch.bfloat16 if is_qwen_image_class(class_name) and torch.cuda.is_available() else torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
variant = "fp16" if torch.cuda.is_available() and has_fp16_variant(model_path) else None
|
||||
if class_name == "QwenImageEditPipeline":
|
||||
from diffusers import QwenImageEditPipeline
|
||||
|
||||
pipe = QwenImageEditPipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
elif class_name == "QwenImagePipeline":
|
||||
from diffusers import QwenImagePipeline
|
||||
|
||||
pipe = QwenImagePipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
elif edit:
|
||||
from diffusers import StableDiffusionXLInpaintPipeline
|
||||
|
||||
pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=dtype,
|
||||
use_safetensors=True,
|
||||
variant=variant,
|
||||
)
|
||||
else:
|
||||
from diffusers import StableDiffusionXLPipeline
|
||||
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=dtype,
|
||||
use_safetensors=True,
|
||||
variant=variant,
|
||||
)
|
||||
return pipe, torch, class_name
|
||||
|
||||
|
||||
def json_response(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
|
||||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
handler.send_response(status)
|
||||
handler.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
handler.send_header("Content-Length", str(len(body)))
|
||||
handler.end_headers()
|
||||
handler.wfile.write(body)
|
||||
|
||||
|
||||
def read_body(handler: BaseHTTPRequestHandler) -> dict:
|
||||
length = int(handler.headers.get("Content-Length") or "0")
|
||||
if length <= 0:
|
||||
return {}
|
||||
return json.loads(handler.rfile.read(length).decode("utf-8"))
|
||||
|
||||
|
||||
def last_user_message(messages: list[dict]) -> str:
|
||||
for message in reversed(messages):
|
||||
if message.get("role") == "user":
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = [part.get("text", "") for part in content if isinstance(part, dict)]
|
||||
return "\n".join(part for part in parts if part)
|
||||
return ""
|
||||
|
||||
|
||||
def load_translation_model() -> None:
|
||||
global LOAD_ERROR, MODEL, TOKENIZER
|
||||
if MODEL is not None and TOKENIZER is not None:
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
model_path = SERVER_CONFIG["model_path"]
|
||||
TOKENIZER = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
MODEL = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else "auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
LOAD_ERROR = None
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
LOAD_ERROR = f"{type(exc).__name__}: {exc}"
|
||||
raise
|
||||
|
||||
|
||||
def load_audio_model() -> None:
|
||||
global LOAD_ERROR, MODEL, PROCESSOR
|
||||
if MODEL is not None and PROCESSOR is not None:
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
||||
|
||||
model_path = SERVER_CONFIG["model_path"]
|
||||
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
MODEL = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
use_safetensors=True,
|
||||
)
|
||||
if torch.cuda.is_available():
|
||||
MODEL.to("cuda:0")
|
||||
PROCESSOR = AutoProcessor.from_pretrained(model_path)
|
||||
SERVER_CONFIG["pipeline"] = pipeline(
|
||||
"automatic-speech-recognition",
|
||||
model=MODEL,
|
||||
tokenizer=PROCESSOR.tokenizer,
|
||||
feature_extractor=PROCESSOR.feature_extractor,
|
||||
torch_dtype=dtype,
|
||||
device=0 if torch.cuda.is_available() else -1,
|
||||
)
|
||||
LOAD_ERROR = None
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
LOAD_ERROR = f"{type(exc).__name__}: {exc}"
|
||||
raise
|
||||
|
||||
|
||||
def load_video_model() -> None:
|
||||
global LOAD_ERROR, MODEL, PROCESSOR
|
||||
if MODEL is not None and PROCESSOR is not None:
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoProcessor
|
||||
|
||||
try:
|
||||
from transformers import Qwen2_5_VLForConditionalGeneration as VisionModel
|
||||
except ImportError:
|
||||
from transformers import AutoModelForVision2Seq as VisionModel
|
||||
|
||||
model_path = SERVER_CONFIG["model_path"]
|
||||
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
MODEL = VisionModel.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=dtype,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
PROCESSOR = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
|
||||
LOAD_ERROR = None
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
LOAD_ERROR = f"{type(exc).__name__}: {exc}"
|
||||
raise
|
||||
|
||||
|
||||
def load_image_model() -> None:
|
||||
global LOAD_ERROR, MODEL
|
||||
if MODEL is not None and MODEL_LOAD_STATE.get("image", {}).get("status") == "loaded":
|
||||
return
|
||||
with IMAGE_MODEL_LOCK:
|
||||
if MODEL is not None and MODEL_LOAD_STATE.get("image", {}).get("status") == "loaded":
|
||||
return
|
||||
if MODEL is not None:
|
||||
set_load_state("image", status="loaded", recovered_at=time.time(), error=None)
|
||||
return
|
||||
started = time.perf_counter()
|
||||
set_load_state("image", status="loading", started_at=time.time(), error=None)
|
||||
try:
|
||||
if not has_active_image_job("edit"):
|
||||
unload_image_edit_model()
|
||||
model_path = SERVER_CONFIG["model_path"]
|
||||
MODEL, torch, class_name = load_diffusers_pipeline(model_path, edit=False)
|
||||
set_load_state("image", status="moving_to_gpu" if torch.cuda.is_available() else "loading", error=None)
|
||||
prepare_diffusers_pipeline(MODEL, torch_module=torch, class_name=class_name)
|
||||
LOAD_ERROR = None
|
||||
set_load_state("image", status="loaded", pipeline_class=class_name, latency_ms=round((time.perf_counter() - started) * 1000), error=None)
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
MODEL = None
|
||||
LOAD_ERROR = f"{type(exc).__name__}: {exc}"
|
||||
set_load_state("image", status="error", error=LOAD_ERROR, latency_ms=round((time.perf_counter() - started) * 1000))
|
||||
raise
|
||||
|
||||
|
||||
def load_image_edit_model() -> None:
|
||||
global LOAD_ERROR, EDIT_MODEL
|
||||
if EDIT_MODEL is not None and MODEL_LOAD_STATE.get("image_edit", {}).get("status") == "loaded":
|
||||
return
|
||||
with IMAGE_EDIT_MODEL_LOCK:
|
||||
if EDIT_MODEL is not None and MODEL_LOAD_STATE.get("image_edit", {}).get("status") == "loaded":
|
||||
return
|
||||
if EDIT_MODEL is not None:
|
||||
set_load_state("image_edit", status="loaded", recovered_at=time.time(), error=None)
|
||||
return
|
||||
started = time.perf_counter()
|
||||
set_load_state("image_edit", status="loading", started_at=time.time(), error=None)
|
||||
try:
|
||||
if not has_active_image_job("generate"):
|
||||
unload_image_model()
|
||||
model_path = SERVER_CONFIG.get("edit_model_path") or SERVER_CONFIG["model_path"]
|
||||
EDIT_MODEL, torch, class_name = load_diffusers_pipeline(model_path, edit=True)
|
||||
set_load_state("image_edit", status="moving_to_gpu" if torch.cuda.is_available() else "loading", error=None)
|
||||
prepare_diffusers_pipeline(EDIT_MODEL, torch_module=torch, class_name=class_name)
|
||||
LOAD_ERROR = None
|
||||
set_load_state("image_edit", status="loaded", pipeline_class=class_name, latency_ms=round((time.perf_counter() - started) * 1000), error=None)
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
EDIT_MODEL = None
|
||||
LOAD_ERROR = f"{type(exc).__name__}: {exc}"
|
||||
set_load_state("image_edit", status="error", error=LOAD_ERROR, latency_ms=round((time.perf_counter() - started) * 1000))
|
||||
raise
|
||||
|
||||
|
||||
def generate_translation(messages: list[dict], max_tokens: int, temperature: float) -> str:
|
||||
import torch
|
||||
|
||||
load_translation_model()
|
||||
assert MODEL is not None
|
||||
assert TOKENIZER is not None
|
||||
|
||||
prompt = last_user_message(messages)
|
||||
if not prompt:
|
||||
prompt = "Translate the input text."
|
||||
system = (
|
||||
"You are a precise translation engine. Preserve formatting, numbers, JSON keys, "
|
||||
"and technical terms unless the user explicitly asks otherwise."
|
||||
)
|
||||
chat = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
|
||||
input_ids = TOKENIZER.apply_chat_template(chat, add_generation_prompt=True, return_tensors="pt").to(MODEL.device)
|
||||
outputs = MODEL.generate(
|
||||
input_ids,
|
||||
max_new_tokens=max_tokens,
|
||||
do_sample=temperature > 0,
|
||||
temperature=max(temperature, 0.01),
|
||||
pad_token_id=TOKENIZER.eos_token_id,
|
||||
)
|
||||
generated = outputs[0][input_ids.shape[-1] :]
|
||||
return TOKENIZER.decode(generated, skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def decode_data_url(value: str, *, field_name: str) -> bytes:
|
||||
data = str(value or "")
|
||||
if "," in data and data.split(",", 1)[0].startswith("data:"):
|
||||
data = data.split(",", 1)[1]
|
||||
if not data:
|
||||
raise ValueError(f"{field_name} is required")
|
||||
return base64.b64decode(data)
|
||||
|
||||
|
||||
def transcribe_audio(payload: dict) -> dict:
|
||||
if SERVER_CONFIG["plugin"] != "audio":
|
||||
raise ValueError("audio transcription is available only for the audio plugin")
|
||||
suffix = Path(str(payload.get("filename") or "audio.wav")).suffix or ".wav"
|
||||
load_audio_model()
|
||||
pipeline_obj = SERVER_CONFIG.get("pipeline")
|
||||
if pipeline_obj is None:
|
||||
raise RuntimeError("audio pipeline is not initialized")
|
||||
|
||||
audio_bytes = decode_data_url(str(payload.get("audio_base64") or ""), field_name="audio_base64")
|
||||
with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as handle:
|
||||
handle.write(audio_bytes)
|
||||
handle.flush()
|
||||
kwargs = {}
|
||||
language = payload.get("language")
|
||||
task = payload.get("task")
|
||||
if language or task:
|
||||
kwargs["generate_kwargs"] = {}
|
||||
if language:
|
||||
kwargs["generate_kwargs"]["language"] = str(language)
|
||||
if task:
|
||||
kwargs["generate_kwargs"]["task"] = str(task)
|
||||
result = pipeline_obj(handle.name, **kwargs)
|
||||
return {
|
||||
"text": str(result.get("text") or "").strip(),
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"bytes": len(audio_bytes),
|
||||
}
|
||||
|
||||
|
||||
def analyze_image(payload: dict) -> dict:
|
||||
if SERVER_CONFIG["plugin"] != "video":
|
||||
raise ValueError("vision analysis is available only for the video plugin")
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
load_video_model()
|
||||
assert MODEL is not None
|
||||
assert PROCESSOR is not None
|
||||
|
||||
prompt = str(payload.get("prompt") or "").strip() or "Опиши изображение и перечисли важные детали."
|
||||
max_tokens = int(payload.get("max_tokens") or 512)
|
||||
image_bytes = decode_data_url(str(payload.get("image_base64") or ""), field_name="image_base64")
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": image},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
text = PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
inputs = PROCESSOR(text=[text], images=[image], return_tensors="pt")
|
||||
device = getattr(MODEL, "device", None)
|
||||
if device is not None:
|
||||
inputs = {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()}
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = MODEL.generate(**inputs, max_new_tokens=max_tokens)
|
||||
input_length = inputs["input_ids"].shape[-1]
|
||||
generated = outputs[0][input_length:]
|
||||
answer = PROCESSOR.decode(generated, skip_special_tokens=True).strip()
|
||||
return {
|
||||
"text": answer,
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"bytes": len(image_bytes),
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
}
|
||||
|
||||
|
||||
def encode_image_data_url(image) -> str:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def clamp_image_dimension(value: object, fallback: int = 1024) -> int:
|
||||
try:
|
||||
number = int(value)
|
||||
except (TypeError, ValueError):
|
||||
number = fallback
|
||||
number = max(512, min(1536, number))
|
||||
return number - (number % 8)
|
||||
|
||||
|
||||
def optional_seed(value: object) -> int | None:
|
||||
if value in {None, ""}:
|
||||
return None
|
||||
return int(value)
|
||||
|
||||
|
||||
def image_generator(seed: int | None):
|
||||
if seed is None:
|
||||
return None
|
||||
import torch
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
return torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
|
||||
def gpu_runtime_status() -> dict:
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return {"available": False, "device_count": 0, "summary": "cuda unavailable"}
|
||||
devices = []
|
||||
for index in range(torch.cuda.device_count()):
|
||||
total = torch.cuda.get_device_properties(index).total_memory
|
||||
allocated = torch.cuda.memory_allocated(index)
|
||||
reserved = torch.cuda.memory_reserved(index)
|
||||
free = max(total - reserved, 0)
|
||||
devices.append(
|
||||
{
|
||||
"index": index,
|
||||
"name": torch.cuda.get_device_name(index),
|
||||
"memory_total_mib": round(total / 1024 / 1024),
|
||||
"memory_allocated_mib": round(allocated / 1024 / 1024),
|
||||
"memory_reserved_mib": round(reserved / 1024 / 1024),
|
||||
"memory_free_mib": round(free / 1024 / 1024),
|
||||
}
|
||||
)
|
||||
summary = " · ".join(
|
||||
f"GPU{item['index']} {item['memory_free_mib']}/{item['memory_total_mib']} MiB free"
|
||||
for item in devices
|
||||
)
|
||||
return {"available": True, "device_count": len(devices), "devices": devices, "summary": summary}
|
||||
except Exception as exc: # pragma: no cover - depends on runtime
|
||||
return {"available": False, "device_count": 0, "summary": str(exc), "error": str(exc)}
|
||||
|
||||
|
||||
def image_job_public(job: dict) -> dict:
|
||||
public = dict(job)
|
||||
payload = public.get("payload")
|
||||
if isinstance(payload, dict):
|
||||
public["payload"] = {
|
||||
key: value
|
||||
for key, value in payload.items()
|
||||
if key not in {"image_base64", "mask_base64"}
|
||||
}
|
||||
return public
|
||||
|
||||
|
||||
def get_image_job(job_id: str) -> dict | None:
|
||||
with IMAGE_JOB_LOCK:
|
||||
job = IMAGE_JOBS.get(job_id)
|
||||
return image_job_public(job) if job else None
|
||||
|
||||
|
||||
def get_image_job_internal(job_id: str) -> dict | None:
|
||||
with IMAGE_JOB_LOCK:
|
||||
job = IMAGE_JOBS.get(job_id)
|
||||
return dict(job) if job else None
|
||||
|
||||
|
||||
def update_image_job(job_id: str, **fields: object) -> dict | None:
|
||||
with IMAGE_JOB_LOCK:
|
||||
job = IMAGE_JOBS.get(job_id)
|
||||
if not job:
|
||||
return None
|
||||
job.update(fields)
|
||||
job["updated_at"] = time.time()
|
||||
return image_job_public(job)
|
||||
|
||||
|
||||
def list_image_jobs(limit: int = 20) -> list[dict]:
|
||||
with IMAGE_JOB_LOCK:
|
||||
jobs = sorted(IMAGE_JOBS.values(), key=lambda item: float(item.get("created_at") or 0), reverse=True)
|
||||
return [image_job_public(job) for job in jobs[:limit]]
|
||||
|
||||
|
||||
def ensure_not_cancelled(cancel_event: threading.Event | None) -> None:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("cancelled")
|
||||
|
||||
|
||||
def diffusers_cancel_callback(cancel_event: threading.Event | None):
|
||||
if cancel_event is None:
|
||||
return None
|
||||
|
||||
def callback(pipe, step, timestep, callback_kwargs):
|
||||
if cancel_event.is_set():
|
||||
setattr(pipe, "_interrupt", True)
|
||||
return callback_kwargs
|
||||
|
||||
return callback
|
||||
|
||||
|
||||
def call_diffusers_pipeline(pipe, *, cancel_event: threading.Event | None = None, **kwargs):
|
||||
callback = diffusers_cancel_callback(cancel_event)
|
||||
if callback is not None:
|
||||
kwargs["callback_on_step_end"] = callback
|
||||
try:
|
||||
return pipe(**kwargs)
|
||||
except TypeError:
|
||||
kwargs.pop("callback_on_step_end", None)
|
||||
ensure_not_cancelled(cancel_event)
|
||||
return pipe(**kwargs)
|
||||
|
||||
|
||||
def generate_image(payload: dict, cancel_event: threading.Event | None = None) -> dict:
|
||||
if SERVER_CONFIG["plugin"] != "image":
|
||||
raise ValueError("image generation is available only for the image plugin")
|
||||
load_image_model()
|
||||
assert MODEL is not None
|
||||
ensure_not_cancelled(cancel_event)
|
||||
|
||||
prompt = str(payload.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
raise ValueError("prompt is required")
|
||||
width = clamp_image_dimension(payload.get("width"), 1024)
|
||||
height = clamp_image_dimension(payload.get("height"), 1024)
|
||||
steps = max(1, min(80, int(payload.get("steps") or 28)))
|
||||
guidance_scale = float(payload.get("guidance_scale") or 6.0)
|
||||
seed = optional_seed(payload.get("seed"))
|
||||
ensure_not_cancelled(cancel_event)
|
||||
class_name = diffusers_class_name(SERVER_CONFIG["model_path"])
|
||||
kwargs = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": str(payload.get("negative_prompt") or "") or None,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_inference_steps": steps,
|
||||
"generator": image_generator(seed),
|
||||
}
|
||||
if is_qwen_image_class(class_name):
|
||||
kwargs["true_cfg_scale"] = guidance_scale
|
||||
else:
|
||||
kwargs["guidance_scale"] = guidance_scale
|
||||
result = call_diffusers_pipeline(MODEL, cancel_event=cancel_event, **kwargs)
|
||||
ensure_not_cancelled(cancel_event)
|
||||
image = result.images[0]
|
||||
return {
|
||||
"image_base64": encode_image_data_url(image),
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
|
||||
def edit_image(payload: dict, cancel_event: threading.Event | None = None) -> dict:
|
||||
if SERVER_CONFIG["plugin"] != "image":
|
||||
raise ValueError("image editing is available only for the image plugin")
|
||||
|
||||
from PIL import Image
|
||||
|
||||
load_image_edit_model()
|
||||
assert EDIT_MODEL is not None
|
||||
ensure_not_cancelled(cancel_event)
|
||||
|
||||
prompt = str(payload.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
raise ValueError("prompt is required")
|
||||
image_bytes = decode_data_url(str(payload.get("image_base64") or ""), field_name="image_base64")
|
||||
width = clamp_image_dimension(payload.get("width"), 1024)
|
||||
height = clamp_image_dimension(payload.get("height"), 1024)
|
||||
steps = max(1, min(80, int(payload.get("steps") or 28)))
|
||||
guidance_scale = float(payload.get("guidance_scale") or 6.0)
|
||||
strength = max(0.0, min(1.0, float(payload.get("strength") or 0.95)))
|
||||
seed = optional_seed(payload.get("seed"))
|
||||
source_image = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((width, height))
|
||||
ensure_not_cancelled(cancel_event)
|
||||
class_name = diffusers_class_name(SERVER_CONFIG.get("edit_model_path") or SERVER_CONFIG["model_path"])
|
||||
kwargs = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": str(payload.get("negative_prompt") or "") or None,
|
||||
"image": source_image,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_inference_steps": steps,
|
||||
"generator": image_generator(seed),
|
||||
}
|
||||
if is_qwen_image_class(class_name):
|
||||
kwargs["true_cfg_scale"] = guidance_scale
|
||||
else:
|
||||
mask_bytes = decode_data_url(str(payload.get("mask_base64") or ""), field_name="mask_base64")
|
||||
kwargs["mask_image"] = Image.open(io.BytesIO(mask_bytes)).convert("L").resize((width, height))
|
||||
kwargs["guidance_scale"] = guidance_scale
|
||||
kwargs["strength"] = strength
|
||||
result = call_diffusers_pipeline(EDIT_MODEL, cancel_event=cancel_event, **kwargs)
|
||||
ensure_not_cancelled(cancel_event)
|
||||
image = result.images[0]
|
||||
return {
|
||||
"image_base64": encode_image_data_url(image),
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
|
||||
def run_image_job(job_id: str) -> None:
|
||||
job = get_image_job_internal(job_id)
|
||||
if not job:
|
||||
return
|
||||
cancel_event = IMAGE_JOB_CANCEL_EVENTS[job_id]
|
||||
operation = str(job.get("operation") or "generate")
|
||||
payload = dict(job.get("payload") or {})
|
||||
started = time.perf_counter()
|
||||
load_key = "image_edit" if operation == "edit" else "image"
|
||||
model_loaded = MODEL_LOAD_STATE.get(load_key, {}).get("status") == "loaded"
|
||||
update_image_job(
|
||||
job_id,
|
||||
status="running" if model_loaded else "loading_model",
|
||||
started_at=time.time(),
|
||||
load_state=MODEL_LOAD_STATE.get(load_key),
|
||||
)
|
||||
try:
|
||||
ensure_not_cancelled(cancel_event)
|
||||
if operation == "edit":
|
||||
load_image_edit_model()
|
||||
update_image_job(job_id, status="running", load_state=MODEL_LOAD_STATE.get(load_key))
|
||||
result = edit_image(payload, cancel_event=cancel_event)
|
||||
else:
|
||||
load_image_model()
|
||||
update_image_job(job_id, status="running", load_state=MODEL_LOAD_STATE.get(load_key))
|
||||
result = generate_image(payload, cancel_event=cancel_event)
|
||||
latency_ms = round((time.perf_counter() - started) * 1000)
|
||||
if cancel_event.is_set():
|
||||
update_image_job(job_id, status="cancelled", latency_ms=latency_ms)
|
||||
else:
|
||||
result["latency_ms"] = latency_ms
|
||||
update_image_job(job_id, status="completed", result=result, latency_ms=latency_ms)
|
||||
except RuntimeError as exc:
|
||||
latency_ms = round((time.perf_counter() - started) * 1000)
|
||||
if str(exc).lower() == "cancelled":
|
||||
update_image_job(job_id, status="cancelled", latency_ms=latency_ms)
|
||||
else:
|
||||
update_image_job(job_id, status="error", error=str(exc), error_type=type(exc).__name__, latency_ms=latency_ms)
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
latency_ms = round((time.perf_counter() - started) * 1000)
|
||||
update_image_job(job_id, status="error", error=str(exc), error_type=type(exc).__name__, latency_ms=latency_ms)
|
||||
|
||||
|
||||
class PluginHandler(BaseHTTPRequestHandler):
|
||||
server_version = "LLMTransformersPlugin/1.0"
|
||||
|
||||
def log_message(self, fmt: str, *args: object) -> None:
|
||||
print(f"{self.address_string()} - {fmt % args}", file=sys.stderr)
|
||||
|
||||
def do_GET(self) -> None:
|
||||
if self.path == "/health":
|
||||
json_response(
|
||||
self,
|
||||
200 if LOAD_ERROR is None else 503,
|
||||
{
|
||||
"status": "ok" if LOAD_ERROR is None else "error",
|
||||
"plugin": SERVER_CONFIG["plugin"],
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"model_path": SERVER_CONFIG["model_path"],
|
||||
"edit_model_path": SERVER_CONFIG.get("edit_model_path"),
|
||||
"loaded": MODEL is not None,
|
||||
"edit_loaded": EDIT_MODEL is not None,
|
||||
"load_state": MODEL_LOAD_STATE,
|
||||
"gpu": gpu_runtime_status(),
|
||||
"error": LOAD_ERROR,
|
||||
},
|
||||
)
|
||||
return
|
||||
if self.path == "/v1/images/jobs":
|
||||
json_response(self, 200, {"jobs": list_image_jobs()})
|
||||
return
|
||||
if self.path.startswith("/v1/images/jobs/"):
|
||||
job_id = self.path.rsplit("/", 1)[-1].split("?", 1)[0]
|
||||
job = get_image_job(job_id)
|
||||
if not job:
|
||||
json_response(self, 404, {"error": {"message": "job not found", "type": "not_found"}})
|
||||
return
|
||||
json_response(self, 200, job)
|
||||
return
|
||||
if self.path == "/v1/models":
|
||||
json_response(
|
||||
self,
|
||||
200,
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": SERVER_CONFIG["served_model_name"],
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "local",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
return
|
||||
json_response(self, 404, {"error": "not found"})
|
||||
|
||||
def do_POST(self) -> None:
|
||||
if self.path == "/v1/images/jobs":
|
||||
try:
|
||||
body = read_body(self)
|
||||
operation = str(body.get("operation") or "generate")
|
||||
if operation not in {"generate", "edit"}:
|
||||
raise ValueError("operation must be generate or edit")
|
||||
payload = body.get("payload") if isinstance(body.get("payload"), dict) else body
|
||||
job_id = str(body.get("job_id") or uuid.uuid4())
|
||||
cancel_event = threading.Event()
|
||||
with IMAGE_JOB_LOCK:
|
||||
IMAGE_JOBS[job_id] = {
|
||||
"id": job_id,
|
||||
"operation": operation,
|
||||
"status": "queued",
|
||||
"created_at": time.time(),
|
||||
"updated_at": time.time(),
|
||||
"payload": dict(payload),
|
||||
}
|
||||
IMAGE_JOB_CANCEL_EVENTS[job_id] = cancel_event
|
||||
thread = threading.Thread(target=run_image_job, args=(job_id,), daemon=True)
|
||||
thread.start()
|
||||
json_response(self, 202, get_image_job(job_id) or {"id": job_id, "status": "queued"})
|
||||
except Exception as exc:
|
||||
json_response(self, 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
|
||||
if self.path.startswith("/v1/images/jobs/") and self.path.endswith("/cancel"):
|
||||
job_id = self.path.removeprefix("/v1/images/jobs/").removesuffix("/cancel").strip("/")
|
||||
event = IMAGE_JOB_CANCEL_EVENTS.get(job_id)
|
||||
if event is None:
|
||||
json_response(self, 404, {"error": {"message": "job not found", "type": "not_found"}})
|
||||
return
|
||||
event.set()
|
||||
job = get_image_job(job_id)
|
||||
if job and job.get("status") in {"queued", "loading_model", "running"}:
|
||||
update_image_job(job_id, status="cancel_requested")
|
||||
json_response(self, 200, get_image_job(job_id) or {"id": job_id, "status": "cancel_requested"})
|
||||
return
|
||||
|
||||
if self.path == "/v1/images/generations":
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
result = generate_image(read_body(self))
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
result["latency_ms"] = round((time.perf_counter() - started) * 1000)
|
||||
json_response(self, 200, result)
|
||||
return
|
||||
|
||||
if self.path == "/v1/images/edits":
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
result = edit_image(read_body(self))
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
result["latency_ms"] = round((time.perf_counter() - started) * 1000)
|
||||
json_response(self, 200, result)
|
||||
return
|
||||
|
||||
if self.path == "/v1/vision/analyze":
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
result = analyze_image(read_body(self))
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
result["latency_ms"] = round((time.perf_counter() - started) * 1000)
|
||||
json_response(self, 200, result)
|
||||
return
|
||||
|
||||
if self.path == "/v1/audio/transcriptions":
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
result = transcribe_audio(read_body(self))
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
result["latency_ms"] = round((time.perf_counter() - started) * 1000)
|
||||
json_response(self, 200, result)
|
||||
return
|
||||
|
||||
if self.path != "/v1/chat/completions":
|
||||
json_response(self, 404, {"error": "not found"})
|
||||
return
|
||||
body = read_body(self)
|
||||
messages = body.get("messages") or []
|
||||
max_tokens = int(body.get("max_tokens") or 512)
|
||||
temperature = float(body.get("temperature") or 0.0)
|
||||
|
||||
if SERVER_CONFIG["plugin"] != "translation":
|
||||
json_response(
|
||||
self,
|
||||
501,
|
||||
{
|
||||
"error": {
|
||||
"message": f"{SERVER_CONFIG['plugin']} chat generation is not implemented in this service yet",
|
||||
"type": "not_implemented",
|
||||
}
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
content = generate_translation(messages, max_tokens=max_tokens, temperature=temperature)
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
json_response(self, 503, {"error": {"message": str(exc), "type": type(exc).__name__}})
|
||||
return
|
||||
|
||||
json_response(
|
||||
self,
|
||||
200,
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": SERVER_CONFIG["served_model_name"],
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": content},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Small OpenAI-like service for local Transformers plugin models.")
|
||||
parser.add_argument("--plugin", required=True, choices=["translation", "audio", "video", "image"])
|
||||
parser.add_argument("--model-path", default=os.environ.get("MODEL_PATH"))
|
||||
parser.add_argument("--edit-model-path", default=os.environ.get("EDIT_MODEL_PATH"))
|
||||
parser.add_argument("--served-model-name", default=os.environ.get("SERVED_MODEL_NAME"))
|
||||
parser.add_argument("--host", default=os.environ.get("HOST", "0.0.0.0"))
|
||||
parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8010")))
|
||||
parser.add_argument("--load-on-start", action="store_true", default=os.environ.get("LOAD_ON_START") == "1")
|
||||
parser.add_argument("--background-load-on-start", action="store_true", default=os.environ.get("BACKGROUND_LOAD_ON_START") == "1")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.model_path:
|
||||
parser.error("--model-path or MODEL_PATH is required")
|
||||
if not args.served_model_name:
|
||||
args.served_model_name = Path(args.model_path).name
|
||||
|
||||
SERVER_CONFIG.update(
|
||||
{
|
||||
"plugin": args.plugin,
|
||||
"model_path": args.model_path,
|
||||
"edit_model_path": args.edit_model_path,
|
||||
"served_model_name": args.served_model_name,
|
||||
}
|
||||
)
|
||||
|
||||
def preload() -> None:
|
||||
try:
|
||||
if args.plugin == "translation":
|
||||
load_translation_model()
|
||||
elif args.plugin == "audio":
|
||||
load_audio_model()
|
||||
elif args.plugin == "video":
|
||||
load_video_model()
|
||||
elif args.plugin == "image":
|
||||
load_image_model()
|
||||
except Exception as exc: # pragma: no cover - depends on host GPU/runtime
|
||||
print(f"preload failed: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
|
||||
|
||||
if args.load_on_start and not args.background_load_on_start:
|
||||
preload()
|
||||
|
||||
server = ThreadingHTTPServer((args.host, args.port), PluginHandler)
|
||||
print(f"{args.plugin} service listening on {args.host}:{args.port} model={args.served_model_name}", flush=True)
|
||||
if args.load_on_start and args.background_load_on_start:
|
||||
threading.Thread(target=preload, daemon=True).start()
|
||||
server.serve_forever()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user