Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import base64 | |
| import io | |
| import json | |
| import os | |
| import secrets | |
| import time | |
| import uuid | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces | |
| import gradio as gr | |
| import httpx | |
| import torch | |
| from fastapi import Request | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| from gradio import Server | |
| from PIL import Image | |
| from peft import PeftConfig, PeftModel | |
| from transformers import AutoProcessor, Gemma4ForConditionalGeneration | |
| BASE_ID = os.environ.get( | |
| "BASE_MODEL_ID", "llmfan46/gemma-4-E4B-it-ultra-uncensored-heretic" | |
| ) | |
| ADAPTER_ID = os.environ.get("ADAPTER_ID", "deathafteryou/gemma-4-hacking") | |
| MODEL_ALIAS = os.environ.get("MODEL_ALIAS", "gemma-4-hacking") | |
| API_KEY = os.environ.get("API_KEY", "") | |
| MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "2048")) | |
| # The adapter ships target_modules as bare leaf names ("q_proj", ...). In this | |
| # architecture the vision and audio projections are Gemma4ClippableLinear | |
| # wrappers around an inner nn.Linear, so a bare leaf matches the wrapper and | |
| # PEFT refuses to inject ("Target module Gemma4ClippableLinear ... is not | |
| # supported"). Match the inner .linear of those wrappers explicitly, and keep | |
| # bare leaves for the text tower, which uses plain nn.Linear. | |
| ADAPTER_TARGET_MODULES = ( | |
| r"(.*\.language_model\..*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))" | |
| r"|(.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)\.linear)" | |
| ) | |
| model = Gemma4ForConditionalGeneration.from_pretrained( | |
| BASE_ID, | |
| dtype=torch.bfloat16, | |
| low_cpu_mem_usage=True, | |
| ) | |
| adapter_config = PeftConfig.from_pretrained(ADAPTER_ID) | |
| adapter_config.target_modules = ADAPTER_TARGET_MODULES | |
| # ZeroGPU reports torch.cuda.is_available() as True outside a @spaces.GPU call, so | |
| # PEFT's infer_device() picks "cuda" and safetensors then tries a real CUDA | |
| # allocation ("No CUDA GPUs are available"). Load the adapter on CPU; the | |
| # following .to("cuda") is handled by ZeroGPU's deferred device aliases. | |
| model = PeftModel.from_pretrained( | |
| model, ADAPTER_ID, config=adapter_config, torch_device="cpu" | |
| ) | |
| model = model.to("cuda").eval() | |
| processor = AutoProcessor.from_pretrained(BASE_ID) | |
| def load_image(source): | |
| if isinstance(source, dict): | |
| source = source.get("url") | |
| if not isinstance(source, str) or not source: | |
| raise ValueError("image source is empty") | |
| if source.startswith("data:"): | |
| _, encoded = source.split(",", 1) | |
| raw = base64.b64decode(encoded) | |
| else: | |
| response = httpx.get(source, follow_redirects=True, timeout=30) | |
| response.raise_for_status() | |
| raw = response.content | |
| return Image.open(io.BytesIO(raw)).convert("RGB") | |
| def normalize_messages(messages): | |
| normalized = [] | |
| images = [] | |
| for message in messages or []: | |
| if not isinstance(message, dict): | |
| continue | |
| role = message.get("role") or "user" | |
| content = message.get("content") | |
| if isinstance(content, str): | |
| normalized.append({"role": role, "content": content}) | |
| continue | |
| if not isinstance(content, list): | |
| normalized.append({"role": role, "content": ""}) | |
| continue | |
| parts = [] | |
| for item in content: | |
| if not isinstance(item, dict): | |
| continue | |
| kind = item.get("type") | |
| if kind == "text": | |
| parts.append({"type": "text", "text": item.get("text") or ""}) | |
| elif kind in ("image", "image_url"): | |
| source = item.get("image_url") or item.get("image") | |
| images.append(load_image(source)) | |
| parts.append({"type": "image"}) | |
| if not parts: | |
| parts = [{"type": "text", "text": ""}] | |
| normalized.append({"role": role, "content": parts}) | |
| return normalized, images | |
| def estimate_duration(messages=None, images=None, temperature=0.0, max_tokens=256, **kwargs): | |
| return max(30, min(120, 20 + int(max_tokens) // 2)) | |
| def generate(messages, images, temperature=0.0, max_tokens=256): | |
| text = processor.apply_chat_template( | |
| messages, add_generation_prompt=True, tokenize=False | |
| ) | |
| arguments = {"text": text, "return_tensors": "pt"} | |
| if images: | |
| arguments["images"] = images | |
| inputs = processor(**arguments) | |
| inputs = { | |
| key: (value.to(model.device) if hasattr(value, "to") else value) | |
| for key, value in inputs.items() | |
| } | |
| with torch.inference_mode(): | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_tokens), | |
| do_sample=float(temperature) > 0, | |
| temperature=max(float(temperature), 1e-5), | |
| use_cache=True, | |
| ) | |
| prompt_tokens = int(inputs["input_ids"].shape[1]) | |
| generated = output[:, prompt_tokens:] | |
| content = processor.batch_decode(generated, skip_special_tokens=True)[0] | |
| return { | |
| "text": content, | |
| "prompt_tokens": prompt_tokens, | |
| "completion_tokens": int(generated.shape[1]), | |
| } | |
| def extract_token(request: Request) -> str: | |
| # X-API-Key takes priority: a private Space consumes the Authorization header | |
| # for the Hub's own owner-only gate, so the API key must be able to travel in | |
| # its own header while Authorization carries the Hugging Face token. | |
| api_key_header = request.headers.get("x-api-key", "").strip() | |
| if api_key_header: | |
| return api_key_header | |
| header = request.headers.get("authorization", "") | |
| if header[:7].lower() == "bearer ": | |
| return header[7:].strip() | |
| return "" | |
| def authorized(request: Request) -> bool: | |
| if not API_KEY: | |
| return False | |
| token = extract_token(request) | |
| if not token: | |
| return False | |
| return secrets.compare_digest(token, API_KEY) | |
| def unauthorized() -> JSONResponse: | |
| return JSONResponse( | |
| { | |
| "error": { | |
| "message": "Invalid or missing API key", | |
| "type": "invalid_request_error", | |
| "code": "unauthorized", | |
| } | |
| }, | |
| status_code=401, | |
| headers={"WWW-Authenticate": "Bearer"}, | |
| ) | |
| app = Server(title="Gemma 4 Hacking Secure API") | |
| def healthz(): | |
| return {"status": "ok", "ready": True, "model": MODEL_ALIAS} | |
| def list_models(request: Request): | |
| if not authorized(request): | |
| return unauthorized() | |
| return { | |
| "object": "list", | |
| "data": [ | |
| { | |
| "id": MODEL_ALIAS, | |
| "object": "model", | |
| "created": 0, | |
| "owned_by": "local", | |
| } | |
| ], | |
| } | |
| def completion_payload(text, prompt_tokens, completion_tokens, model_name): | |
| return { | |
| "id": f"chatcmpl-{uuid.uuid4().hex}", | |
| "object": "chat.completion", | |
| "created": int(time.time()), | |
| "model": model_name, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "message": {"role": "assistant", "content": text}, | |
| "finish_reason": "stop", | |
| } | |
| ], | |
| "usage": { | |
| "prompt_tokens": prompt_tokens, | |
| "completion_tokens": completion_tokens, | |
| "total_tokens": prompt_tokens + completion_tokens, | |
| }, | |
| } | |
| def stream_chunks(payload): | |
| choice = payload["choices"][0] | |
| content = choice["message"]["content"] | |
| base = { | |
| "id": payload["id"], | |
| "object": "chat.completion.chunk", | |
| "created": payload["created"], | |
| "model": payload["model"], | |
| } | |
| for index in range(0, len(content), 24): | |
| chunk = dict(base) | |
| chunk["choices"] = [ | |
| { | |
| "index": 0, | |
| "delta": {"content": content[index:index + 24]}, | |
| "finish_reason": None, | |
| } | |
| ] | |
| yield f"data: {json.dumps(chunk)}\n\n" | |
| final = dict(base) | |
| final["choices"] = [{"index": 0, "delta": {}, "finish_reason": "stop"}] | |
| yield f"data: {json.dumps(final)}\n\n" | |
| yield "data: [DONE]\n\n" | |
| def chat_completions(request: Request, payload: dict): | |
| if not authorized(request): | |
| return unauthorized() | |
| try: | |
| messages, images = normalize_messages(payload.get("messages")) | |
| except Exception as exc: | |
| return JSONResponse( | |
| {"error": {"message": str(exc), "type": "invalid_request_error"}}, | |
| status_code=400, | |
| ) | |
| max_tokens = payload.get("max_completion_tokens") or payload.get("max_tokens") or 256 | |
| max_tokens = max(1, min(int(max_tokens), MAX_NEW_TOKENS)) | |
| temperature = float(payload.get("temperature") or 0.0) | |
| try: | |
| result = generate(messages, images, temperature, max_tokens) | |
| except Exception as exc: | |
| return JSONResponse( | |
| {"error": {"message": str(exc), "type": "server_error"}}, | |
| status_code=500, | |
| ) | |
| body = completion_payload( | |
| result["text"], | |
| result["prompt_tokens"], | |
| result["completion_tokens"], | |
| payload.get("model") or MODEL_ALIAS, | |
| ) | |
| if payload.get("stream"): | |
| return StreamingResponse(stream_chunks(body), media_type="text/event-stream") | |
| return JSONResponse(body) | |
| def completions(request: Request, payload: dict): | |
| if not authorized(request): | |
| return unauthorized() | |
| prompt = payload.get("prompt") or "" | |
| if isinstance(prompt, list): | |
| prompt = "\n".join(str(item) for item in prompt) | |
| max_tokens = payload.get("max_tokens") or 256 | |
| max_tokens = max(1, min(int(max_tokens), MAX_NEW_TOKENS)) | |
| temperature = float(payload.get("temperature") or 0.0) | |
| try: | |
| result = generate( | |
| [{"role": "user", "content": str(prompt)}], [], temperature, max_tokens | |
| ) | |
| except Exception as exc: | |
| return JSONResponse( | |
| {"error": {"message": str(exc), "type": "server_error"}}, | |
| status_code=500, | |
| ) | |
| return JSONResponse( | |
| { | |
| "id": f"cmpl-{uuid.uuid4().hex}", | |
| "object": "text_completion", | |
| "created": int(time.time()), | |
| "model": payload.get("model") or MODEL_ALIAS, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "text": result["text"], | |
| "finish_reason": "stop", | |
| } | |
| ], | |
| "usage": { | |
| "prompt_tokens": result["prompt_tokens"], | |
| "completion_tokens": result["completion_tokens"], | |
| "total_tokens": result["prompt_tokens"] + result["completion_tokens"], | |
| }, | |
| } | |
| ) | |
| with gr.Blocks(title="Gemma 4 Hacking Secure API") as demo: | |
| gr.Markdown( | |
| "# Gemma 4 Hacking Secure API\n\n" | |
| "`deathafteryou/gemma-4-hacking` served on ZeroGPU as an " | |
| "OpenAI-compatible API.\n\n" | |
| "All `/v1` routes require `Authorization: Bearer <API_KEY>` or " | |
| "`X-API-Key: <API_KEY>`. `/healthz` is the only public route.\n\n" | |
| "`POST /v1/chat/completions` accepts text and image content parts. " | |
| "Use the `openai` Python client with this Space URL plus `/v1`." | |
| ) | |
| gr.Markdown( | |
| "The API key is stored in the Space secret `API_KEY`. Keep the Space " | |
| "URL and key private." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(show_error=True, _app=app, ssr_mode=False) | |