931 lines
36 KiB
Python
931 lines
36 KiB
Python
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())
|