| from __future__ import annotations |
|
|
| import base64 |
| import json |
| import hashlib |
| import os |
| import re |
| import tempfile |
| import threading |
| import time |
| import traceback |
| import types |
| import urllib.parse |
| from datetime import datetime, timezone |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from peft import PeftModel, get_peft_model_state_dict |
| from peft.tuners.lora.layer import LoraLayer |
| from safetensors import safe_open |
| from transformers import AutoModelForMultimodalLM, AutoProcessor |
| from transformers.video_utils import load_video |
|
|
|
|
| MODEL_DIR = Path(os.environ["GEMMA4_MODEL_DIR"]) |
| ADAPTER_DIR = Path(os.environ["GEMMA4_ADAPTER_DIR"]) |
| TEMPLATE_PATH = Path(os.environ["GEMMA4_TEMPLATE_PATH"]) |
| STATE_PATH = Path(os.environ["GEMMA4_STATE_PATH"]) |
| HOST = os.environ.get("GEMMA4_HOST", "127.0.0.1") |
| PORT = int(os.environ.get("GEMMA4_PORT", "8091")) |
| ALIAS = os.environ.get( |
| "GEMMA4_ALIAS", "anru-human10mb-selected-native-multimodal-128k-mix045" |
| ) |
| LORA_SCALE = float(os.environ.get("GEMMA4_LORA_SCALE", "0.45")) |
| MAX_CONTEXT = int(os.environ.get("GEMMA4_MAX_CONTEXT", "131072")) |
| MAX_IMAGE_TOKENS = int(os.environ.get("GEMMA4_MAX_IMAGE_TOKENS", "1120")) |
| MAX_GPU_MEMORY = os.environ.get("GEMMA4_MAX_GPU_MEMORY", "14GiB") |
| MAX_CPU_MEMORY = os.environ.get("GEMMA4_MAX_CPU_MEMORY", "48GiB") |
| OFFLOAD_DIR = Path(os.environ["GEMMA4_OFFLOAD_DIR"]) |
| ATTN_IMPLEMENTATION = os.environ.get("GEMMA4_ATTN_IMPLEMENTATION", "sdpa") |
| SYSTEM_PROMPT_PATH = Path( |
| os.environ.get("GEMMA4_SYSTEM_PROMPT_PATH", str(Path(__file__).with_name("system_prompt.md"))) |
| ) |
|
|
|
|
| for required in (MODEL_DIR, ADAPTER_DIR, TEMPLATE_PATH, SYSTEM_PROMPT_PATH): |
| if not required.exists(): |
| raise FileNotFoundError(required) |
| OFFLOAD_DIR.mkdir(parents=True, exist_ok=True) |
| STATE_PATH.parent.mkdir(parents=True, exist_ok=True) |
| DEFAULT_SYSTEM_PROMPT = SYSTEM_PROMPT_PATH.read_text(encoding="utf-8-sig").strip() |
| if not DEFAULT_SYSTEM_PROMPT: |
| raise ValueError(f"system prompt is empty: {SYSTEM_PROMPT_PATH}") |
| SYSTEM_PROMPT_BYTES = SYSTEM_PROMPT_PATH.read_bytes() |
| SYSTEM_PROMPT_SHA256 = hashlib.sha256(SYSTEM_PROMPT_BYTES).hexdigest().upper() |
| SYSTEM_PROMPT_LINE_COUNT = len(DEFAULT_SYSTEM_PROMPT.splitlines()) |
|
|
| started_at = time.perf_counter() |
| print(f"[{datetime.now(timezone.utc).isoformat()}] loading processor", flush=True) |
| processor = AutoProcessor.from_pretrained(MODEL_DIR, local_files_only=True) |
| processor.chat_template = TEMPLATE_PATH.read_text(encoding="utf-8") |
|
|
|
|
| def _fetch_videos_pyav(self, video_url_or_urls, sample_indices_fn=None): |
| if isinstance(video_url_or_urls, list): |
| return list( |
| zip(*[_fetch_videos_pyav(self, item, sample_indices_fn=sample_indices_fn) for item in video_url_or_urls]) |
| ) |
| return load_video(video_url_or_urls, backend="pyav", sample_indices_fn=sample_indices_fn) |
|
|
|
|
| if getattr(processor, "video_processor", None) is not None: |
| processor.video_processor.fetch_videos = types.MethodType(_fetch_videos_pyav, processor.video_processor) |
|
|
| print(f"[{datetime.now(timezone.utc).isoformat()}] loading Gemma4 multimodal weights", flush=True) |
| model = AutoModelForMultimodalLM.from_pretrained( |
| MODEL_DIR, |
| dtype=torch.bfloat16, |
| device_map="auto", |
| max_memory={0: MAX_GPU_MEMORY, "cpu": MAX_CPU_MEMORY}, |
| offload_folder=str(OFFLOAD_DIR), |
| offload_buffers=True, |
| low_cpu_mem_usage=True, |
| local_files_only=True, |
| attn_implementation=ATTN_IMPLEMENTATION, |
| ) |
|
|
| print(f"[{datetime.now(timezone.utc).isoformat()}] attaching LoRA adapter", flush=True) |
| model = PeftModel.from_pretrained( |
| model, |
| ADAPTER_DIR, |
| is_trainable=False, |
| local_files_only=True, |
| key_mapping={r"^model\.": "model.language_model."}, |
| device_map="auto", |
| max_memory={0: MAX_GPU_MEMORY, "cpu": MAX_CPU_MEMORY}, |
| offload_folder=str(OFFLOAD_DIR), |
| ) |
|
|
| lora_modules = 0 |
| for module in model.modules(): |
| if isinstance(module, LoraLayer) and "default" in module.active_adapters: |
| module.scale_layer(LORA_SCALE) |
| lora_modules += 1 |
| if not lora_modules: |
| raise RuntimeError("LoRA adapter attached zero modules") |
|
|
| adapter_state = get_peft_model_state_dict(model, adapter_name="default") |
| adapter_tensor_count = len(adapter_state) |
| adapter_file = ADAPTER_DIR / "adapter_model.safetensors" |
| with safe_open(adapter_file, framework="pt", device="cpu") as adapter_reader: |
| source_adapter_tensor_count = len(adapter_reader.keys()) |
| adapter_nonzero_tensors = sum( |
| int(torch.count_nonzero(adapter_reader.get_tensor(key)).item() > 0) |
| for key in adapter_reader.keys() |
| ) |
| if adapter_tensor_count == 0 or adapter_nonzero_tensors == 0: |
| raise RuntimeError("LoRA adapter state is empty") |
| if adapter_tensor_count != source_adapter_tensor_count: |
| raise RuntimeError( |
| f"LoRA tensor count mismatch: runtime={adapter_tensor_count}, source={source_adapter_tensor_count}" |
| ) |
|
|
| model.eval() |
| generation_lock = threading.Lock() |
| loaded_seconds = round(time.perf_counter() - started_at, 3) |
|
|
|
|
| def _json_bytes(value: Any) -> bytes: |
| return json.dumps(value, ensure_ascii=False).encode("utf-8") |
|
|
|
|
| def _data_url_to_file(url: str, temp_paths: list[Path]) -> str: |
| match = re.fullmatch(r"data:([^;,]+)?(?:;charset=[^;,]+)?;base64,(.+)", url, re.DOTALL) |
| if not match: |
| return url |
| mime = match.group(1) or "application/octet-stream" |
| suffixes = { |
| "image/jpeg": ".jpg", |
| "image/png": ".png", |
| "image/webp": ".webp", |
| "audio/wav": ".wav", |
| "audio/mpeg": ".mp3", |
| "video/mp4": ".mp4", |
| } |
| suffix = suffixes.get(mime, ".bin") |
| fd, name = tempfile.mkstemp(prefix="gemma4-input-", suffix=suffix) |
| os.close(fd) |
| path = Path(name) |
| path.write_bytes(base64.b64decode(match.group(2), validate=True)) |
| temp_paths.append(path) |
| return str(path) |
|
|
|
|
| def _normalise_url(value: Any, temp_paths: list[Path]) -> str: |
| if isinstance(value, dict): |
| value = value.get("url") |
| if not isinstance(value, str) or not value: |
| raise ValueError("multimodal content is missing its URL/data") |
| if value.startswith("data:"): |
| return _data_url_to_file(value, temp_paths) |
| if value.startswith("file://"): |
| return urllib.parse.unquote(urllib.parse.urlparse(value).path.lstrip("/") if os.name == "nt" else urllib.parse.urlparse(value).path) |
| return value |
|
|
|
|
| def _normalise_messages( |
| messages: Any, temp_paths: list[Path], use_default_system_prompt: bool |
| ) -> list[dict[str, Any]]: |
| if not isinstance(messages, list) or not messages: |
| raise ValueError("messages must be a non-empty list") |
| result: list[dict[str, Any]] = [] |
| for message in messages: |
| role = message.get("role") |
| content = message.get("content", "") |
| if role not in {"system", "user", "assistant"}: |
| raise ValueError(f"unsupported role: {role!r}") |
| if isinstance(content, str): |
| result.append({"role": role, "content": content}) |
| continue |
| if not isinstance(content, list): |
| raise ValueError("message content must be text or a content-part list") |
| parts: list[dict[str, Any]] = [] |
| for part in content: |
| kind = part.get("type") |
| if kind in {"text", "input_text"}: |
| parts.append({"type": "text", "text": str(part.get("text", ""))}) |
| elif kind in {"image", "image_url", "input_image"}: |
| source = part.get("url", part.get("image_url", part.get("image"))) |
| parts.append({"type": "image", "url": _normalise_url(source, temp_paths)}) |
| elif kind in {"audio", "audio_url", "input_audio"}: |
| source = part.get("url", part.get("audio_url", part.get("audio"))) |
| if source is None and isinstance(part.get("input_audio"), dict): |
| audio = part["input_audio"] |
| fmt = audio.get("format", "wav") |
| source = f"data:audio/{fmt};base64,{audio.get('data', '')}" |
| parts.append({"type": "audio", "url": _normalise_url(source, temp_paths)}) |
| elif kind in {"video", "video_url", "input_video"}: |
| source = part.get("url", part.get("video_url", part.get("video"))) |
| parts.append({"type": "video", "url": _normalise_url(source, temp_paths)}) |
| else: |
| raise ValueError(f"unsupported content part: {kind!r}") |
| result.append({"role": role, "content": parts}) |
| if use_default_system_prompt: |
| if result and result[0]["role"] == "system": |
| existing = result[0]["content"] |
| if isinstance(existing, str): |
| merged = DEFAULT_SYSTEM_PROMPT + "\n\n---\n\n" + existing |
| else: |
| merged = [{"type": "text", "text": DEFAULT_SYSTEM_PROMPT}, *existing] |
| result[0] = {"role": "system", "content": merged} |
| else: |
| result.insert(0, {"role": "system", "content": DEFAULT_SYSTEM_PROMPT}) |
| return result |
|
|
|
|
| def _complete(payload: dict[str, Any]) -> dict[str, Any]: |
| temp_paths: list[Path] = [] |
| began = time.perf_counter() |
| try: |
| use_default_system_prompt = payload.get("use_default_system_prompt", True) |
| if not isinstance(use_default_system_prompt, bool): |
| raise ValueError("use_default_system_prompt must be a JSON boolean") |
| messages = _normalise_messages( |
| payload.get("messages"), temp_paths, use_default_system_prompt |
| ) |
| enable_thinking = bool(payload.get("enable_thinking", False)) |
| inputs = processor.apply_chat_template( |
| messages, |
| tokenize=True, |
| return_dict=True, |
| return_tensors="pt", |
| add_generation_prompt=True, |
| enable_thinking=enable_thinking, |
| processor_kwargs={"images_kwargs": {"max_soft_tokens": MAX_IMAGE_TOKENS}}, |
| ) |
| prompt_tokens = int(inputs["input_ids"].shape[-1]) |
| max_new_tokens = max(1, min(int(payload.get("max_tokens", payload.get("max_completion_tokens", 256))), 4096)) |
| if prompt_tokens + max_new_tokens > MAX_CONTEXT: |
| raise ValueError( |
| f"requested {prompt_tokens + max_new_tokens} tokens exceeds {MAX_CONTEXT}-token service context" |
| ) |
| temperature = float(payload.get("temperature", 0.0)) |
| top_p = float(payload.get("top_p", 0.95)) |
| inputs = inputs.to(model.device) |
| generate_args: dict[str, Any] = { |
| "max_new_tokens": max_new_tokens, |
| "do_sample": temperature > 0.0, |
| "use_cache": True, |
| } |
| if temperature > 0.0: |
| generate_args.update(temperature=temperature, top_p=top_p) |
| with generation_lock, torch.inference_mode(): |
| generated = model.generate(**inputs, **generate_args) |
| completion_tokens = int(generated.shape[-1] - prompt_tokens) |
| text = processor.decode(generated[0][prompt_tokens:], skip_special_tokens=True) |
| for stop in payload.get("stop", []) if isinstance(payload.get("stop", []), list) else [payload.get("stop")]: |
| if stop and stop in text: |
| text = text.split(stop, 1)[0] |
| elapsed = round(time.perf_counter() - began, 3) |
| return { |
| "id": f"chatcmpl-native-{int(time.time() * 1000)}", |
| "object": "chat.completion", |
| "created": int(time.time()), |
| "model": ALIAS, |
| "choices": [ |
| { |
| "index": 0, |
| "message": {"role": "assistant", "content": text}, |
| "finish_reason": "stop" if completion_tokens < max_new_tokens else "length", |
| } |
| ], |
| "usage": { |
| "prompt_tokens": prompt_tokens, |
| "completion_tokens": completion_tokens, |
| "total_tokens": prompt_tokens + completion_tokens, |
| }, |
| "timing": {"elapsed_seconds": elapsed}, |
| "system_prompt": { |
| "default_injected": use_default_system_prompt, |
| "sha256": SYSTEM_PROMPT_SHA256 if use_default_system_prompt else None, |
| }, |
| } |
| finally: |
| for path in temp_paths: |
| try: |
| path.unlink() |
| except OSError: |
| pass |
|
|
|
|
| health = { |
| "status": "ok", |
| "backend": "transformers-native", |
| "model": ALIAS, |
| "modalities": ["text", "image", "audio", "video"], |
| "context_tokens": MAX_CONTEXT, |
| "attention_implementation": ATTN_IMPLEMENTATION, |
| "default_system_prompt": True, |
| "system_prompt_file": SYSTEM_PROMPT_PATH.name, |
| "system_prompt_sha256": SYSTEM_PROMPT_SHA256, |
| "system_prompt_bytes": len(SYSTEM_PROMPT_BYTES), |
| "system_prompt_chars": len(DEFAULT_SYSTEM_PROMPT), |
| "system_prompt_lines": SYSTEM_PROMPT_LINE_COUNT, |
| "image_max_soft_tokens": MAX_IMAGE_TOKENS, |
| "lora_scale": LORA_SCALE, |
| "lora_modules": lora_modules, |
| "adapter_tensor_count": adapter_tensor_count, |
| "adapter_nonzero_tensors": adapter_nonzero_tensors, |
| "load_seconds": loaded_seconds, |
| "cuda_available": torch.cuda.is_available(), |
| "cuda_allocated_bytes": torch.cuda.memory_allocated() if torch.cuda.is_available() else 0, |
| "device_map": getattr(model, "hf_device_map", None), |
| } |
| STATE_PATH.write_text(json.dumps(health, ensure_ascii=False, indent=2), encoding="utf-8") |
| print(json.dumps(health, ensure_ascii=False), flush=True) |
|
|
|
|
| class Handler(BaseHTTPRequestHandler): |
| server_version = "Gemma4Native/1.0" |
|
|
| def _send(self, status: int, value: Any) -> None: |
| body = _json_bytes(value) |
| self.send_response(status) |
| self.send_header("Content-Type", "application/json; charset=utf-8") |
| self.send_header("Content-Length", str(len(body))) |
| self.end_headers() |
| self.wfile.write(body) |
|
|
| def do_GET(self) -> None: |
| path = urllib.parse.urlparse(self.path).path |
| if path == "/health": |
| self._send(200, health) |
| elif path == "/v1/models": |
| self._send(200, {"object": "list", "data": [{"id": ALIAS, "object": "model", "owned_by": "local"}]}) |
| elif path == "/props": |
| self._send( |
| 200, |
| { |
| "model_alias": ALIAS, |
| "total_slots": 1, |
| "default_generation_settings": {"n_ctx": MAX_CONTEXT}, |
| "modalities": health["modalities"], |
| "backend": health["backend"], |
| }, |
| ) |
| elif path == "/lora-adapters": |
| self._send(200, [{"id": 0, "scale": LORA_SCALE, "path": str(ADAPTER_DIR)}]) |
| else: |
| self._send(404, {"error": {"message": "not found", "type": "invalid_request_error"}}) |
|
|
| def do_POST(self) -> None: |
| path = urllib.parse.urlparse(self.path).path |
| if path != "/v1/chat/completions": |
| self._send(404, {"error": {"message": "not found", "type": "invalid_request_error"}}) |
| return |
| try: |
| length = int(self.headers.get("Content-Length", "0")) |
| payload = json.loads(self.rfile.read(length)) |
| if payload.get("stream"): |
| raise ValueError("streaming is not enabled on this local backend") |
| self._send(200, _complete(payload)) |
| except ValueError as exc: |
| self._send(400, {"error": {"message": str(exc), "type": "invalid_request_error"}}) |
| except Exception as exc: |
| traceback.print_exc() |
| self._send(500, {"error": {"message": repr(exc), "type": "server_error"}}) |
|
|
| def log_message(self, fmt: str, *args: Any) -> None: |
| print(f"[{datetime.now(timezone.utc).isoformat()}] {self.client_address[0]} {fmt % args}", flush=True) |
|
|
|
|
| print(f"[{datetime.now(timezone.utc).isoformat()}] listening on http://{HOST}:{PORT}", flush=True) |
| ThreadingHTTPServer((HOST, PORT), Handler).serve_forever() |
|
|
