462 lines
15 KiB
Python
462 lines
15 KiB
Python
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.10,<3.11"
|
|
# dependencies = [
|
|
# "torch==2.9.1+cu128",
|
|
# "torchvision==0.24.1+cu128",
|
|
# "torchaudio==2.9.1+cu128",
|
|
# "diffusers @ git+https://github.com/huggingface/diffusers.git",
|
|
# "transformers>=4.57.0",
|
|
# "accelerate>=1.0.0",
|
|
# "safetensors>=0.4.0",
|
|
# "huggingface_hub>=0.25.0",
|
|
# "numpy>=1.26.0",
|
|
# "pillow>=10.0.0",
|
|
# "fastapi>=0.115.0",
|
|
# "uvicorn>=0.29.0",
|
|
# ]
|
|
# [tool.uv]
|
|
# extra-index-url = ["https://download.pytorch.org/whl/cu128"]
|
|
# ///
|
|
|
|
import argparse
|
|
import base64
|
|
import binascii
|
|
import io
|
|
import os
|
|
import random
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
|
|
|
|
import torch
|
|
from diffusers import Flux2KleinPipeline
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.responses import HTMLResponse, Response
|
|
from pydantic import BaseModel, Field
|
|
from PIL import Image
|
|
|
|
|
|
def _resolve_model_size() -> str:
|
|
parser = argparse.ArgumentParser(add_help=False)
|
|
parser.add_argument("--model", choices=["4b", "9b"], default="9b")
|
|
args, _ = parser.parse_known_args()
|
|
return args.model
|
|
|
|
|
|
def _checkpoint_for(model_size: str) -> str:
|
|
if model_size == "4b":
|
|
return "black-forest-labs/FLUX.2-klein-4B"
|
|
return "black-forest-labs/FLUX.2-klein-9B"
|
|
|
|
|
|
MODEL_SIZE = _resolve_model_size()
|
|
CHECKPOINT = _checkpoint_for(MODEL_SIZE)
|
|
USE_COMPILE = os.environ.get("FLUX2_COMPILE", "1") != "0"
|
|
WARMUP_RUNS = int(os.environ.get("FLUX2_WARMUP_RUNS", "2"))
|
|
WARMUP_STEPS = int(os.environ.get("FLUX2_WARMUP_STEPS", "4"))
|
|
WARMUP_PROMPT = "A high-resolution photo of a red fox in natural daylight"
|
|
MAX_SEED = (2**31) - 1
|
|
MAX_PIXELS = 4 * 1024 * 1024
|
|
|
|
DEFAULT_STEPS = 4
|
|
DEFAULT_GUIDANCE = 1.0
|
|
|
|
app = FastAPI(title="FLUX2 Klein Private Server")
|
|
PIPE: Flux2KleinPipeline | None = None
|
|
PIPE_LOCK = threading.Lock()
|
|
|
|
|
|
class GenerateRequest(BaseModel):
|
|
prompt: str = Field(..., min_length=1)
|
|
width: int = 1024
|
|
height: int = 1024
|
|
num_inference_steps: int = DEFAULT_STEPS
|
|
guidance_scale: float = DEFAULT_GUIDANCE
|
|
seed: int = 0
|
|
randomize_seed: bool = False
|
|
image_data_url: str | None = None
|
|
|
|
|
|
class CompatRequest(BaseModel):
|
|
model: str | None = None
|
|
prompt: str = Field(..., min_length=1)
|
|
imageDataUrl: str | None = None
|
|
size: str | None = None
|
|
width: int | None = None
|
|
height: int | None = None
|
|
num_inference_steps: int | None = None
|
|
guidance_scale: float | None = None
|
|
seed: int | None = None
|
|
randomize_seed: bool = False
|
|
response_format: str = "b64_json"
|
|
n: int = 1
|
|
negative_prompt: str | None = None
|
|
negativePrompt: str | None = None
|
|
|
|
|
|
def _ensure_cuda() -> None:
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required for this server")
|
|
|
|
|
|
def _compile_pipe(pipe: Flux2KleinPipeline) -> bool:
|
|
if not USE_COMPILE:
|
|
return False
|
|
print("Compiling transformer and VAE decoder...")
|
|
start = time.perf_counter()
|
|
pipe.transformer = torch.compile(
|
|
pipe.transformer,
|
|
mode="max-autotune",
|
|
dynamic=False,
|
|
fullgraph=True,
|
|
)
|
|
pipe.vae.decode = torch.compile(
|
|
pipe.vae.decode,
|
|
mode="max-autotune",
|
|
dynamic=False,
|
|
fullgraph=True,
|
|
)
|
|
print(f"Compile finished in {time.perf_counter() - start:.2f}s")
|
|
return True
|
|
|
|
|
|
def _warmup_pipe(pipe: Flux2KleinPipeline) -> None:
|
|
if WARMUP_RUNS <= 0:
|
|
return
|
|
print(f"Warmup started ({WARMUP_RUNS} runs)...")
|
|
start = time.perf_counter()
|
|
for run_idx in range(WARMUP_RUNS):
|
|
generator = torch.Generator(device="cuda").manual_seed(run_idx)
|
|
pipe(
|
|
prompt=WARMUP_PROMPT,
|
|
height=1024,
|
|
width=1024,
|
|
num_inference_steps=WARMUP_STEPS,
|
|
guidance_scale=1.0,
|
|
generator=generator,
|
|
).images[0]
|
|
torch.cuda.synchronize()
|
|
print(f"Warmup finished in {time.perf_counter() - start:.2f}s")
|
|
|
|
|
|
def _load_pipe(ckpt: str) -> Flux2KleinPipeline:
|
|
print(f"Loading checkpoint: {ckpt}")
|
|
pipe = Flux2KleinPipeline.from_pretrained(ckpt, torch_dtype=torch.bfloat16)
|
|
pipe = pipe.to("cuda")
|
|
pipe.transformer.fuse_qkv_projections()
|
|
pipe.vae.fuse_qkv_projections()
|
|
pipe.vae.to(memory_format=torch.channels_last)
|
|
pipe.transformer.set_attention_backend("_native_flash")
|
|
if _compile_pipe(pipe):
|
|
_warmup_pipe(pipe)
|
|
return pipe
|
|
|
|
|
|
def _normalize_dims(width: int, height: int) -> tuple[int, int]:
|
|
if width <= 0 or height <= 0:
|
|
raise HTTPException(status_code=422, detail="width and height must be > 0")
|
|
if width % 16 != 0 or height % 16 != 0:
|
|
raise HTTPException(
|
|
status_code=422, detail="width and height must be multiples of 16"
|
|
)
|
|
if width * height > MAX_PIXELS:
|
|
raise HTTPException(status_code=422, detail="max resolution is 4 megapixels")
|
|
return width, height
|
|
|
|
|
|
def _parse_size(
|
|
size: str | None, width: int | None, height: int | None
|
|
) -> tuple[int, int]:
|
|
if width and height:
|
|
return _normalize_dims(width, height)
|
|
if not size:
|
|
return _normalize_dims(1024, 1024)
|
|
parts = size.lower().split("x")
|
|
if len(parts) != 2:
|
|
raise HTTPException(status_code=422, detail="size must look like 1024x1024")
|
|
try:
|
|
parsed_w = int(parts[0])
|
|
parsed_h = int(parts[1])
|
|
except ValueError as exc:
|
|
raise HTTPException(
|
|
status_code=422, detail="size values must be integers"
|
|
) from exc
|
|
return _normalize_dims(parsed_w, parsed_h)
|
|
|
|
|
|
def _decode_image_data_url(image_data_url: str | None) -> Image.Image | None:
|
|
if not image_data_url:
|
|
return None
|
|
if "," not in image_data_url:
|
|
raise HTTPException(status_code=422, detail="imageDataUrl must be a data URL")
|
|
_, payload = image_data_url.split(",", 1)
|
|
try:
|
|
raw = base64.b64decode(payload)
|
|
except (ValueError, binascii.Error) as exc:
|
|
raise HTTPException(
|
|
status_code=422, detail="invalid base64 in imageDataUrl"
|
|
) from exc
|
|
with Image.open(io.BytesIO(raw)) as image:
|
|
return image.convert("RGB")
|
|
|
|
|
|
def _generate_png_bytes(
|
|
prompt: str,
|
|
width: int,
|
|
height: int,
|
|
num_inference_steps: int,
|
|
guidance_scale: float,
|
|
seed: int,
|
|
image: Image.Image | None,
|
|
) -> tuple[bytes, int]:
|
|
if PIPE is None:
|
|
raise HTTPException(status_code=503, detail="model not loaded yet")
|
|
|
|
generator = torch.Generator(device="cuda").manual_seed(seed)
|
|
kwargs: dict[str, Any] = {
|
|
"prompt": prompt,
|
|
"height": height,
|
|
"width": width,
|
|
"num_inference_steps": num_inference_steps,
|
|
"guidance_scale": guidance_scale,
|
|
"generator": generator,
|
|
}
|
|
|
|
with PIPE_LOCK:
|
|
try:
|
|
if image is not None:
|
|
try:
|
|
output = PIPE(image=image, **kwargs)
|
|
except TypeError:
|
|
output = PIPE(images=[image], **kwargs)
|
|
else:
|
|
output = PIPE(**kwargs)
|
|
except Exception as exc:
|
|
raise HTTPException(
|
|
status_code=500, detail=f"generation failed: {exc}"
|
|
) from exc
|
|
|
|
result_image = output.images[0]
|
|
buf = io.BytesIO()
|
|
result_image.save(buf, format="PNG")
|
|
return buf.getvalue(), seed
|
|
|
|
|
|
@app.on_event("startup")
|
|
def _startup() -> None:
|
|
global PIPE
|
|
_ensure_cuda()
|
|
PIPE = _load_pipe(CHECKPOINT)
|
|
print("Server ready")
|
|
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
def index() -> str:
|
|
return """<!doctype html>
|
|
<html lang=\"en\">
|
|
<head>
|
|
<meta charset=\"utf-8\" />
|
|
<meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />
|
|
<title>FLUX.2 Klein Playground</title>
|
|
<style>
|
|
:root { --bg:#f4f0e8; --card:#fff9ee; --text:#1f1f1f; --accent:#125b50; --muted:#6b6b6b; }
|
|
* { box-sizing: border-box; }
|
|
body { margin:0; font-family: ui-sans-serif, -apple-system, Segoe UI, sans-serif; background: radial-gradient(circle at top left,#fff6df,var(--bg)); color:var(--text); }
|
|
.wrap { max-width: 980px; margin: 1.5rem auto; padding: 1rem; }
|
|
.card { background: var(--card); border: 1px solid #eadfc8; border-radius: 14px; padding: 1rem; box-shadow: 0 8px 20px rgba(0,0,0,.06); }
|
|
h1 { margin: 0 0 .8rem 0; font-size: 1.4rem; }
|
|
.grid { display:grid; gap:.8rem; grid-template-columns: 1fr 1fr; }
|
|
.full { grid-column: 1 / -1; }
|
|
label { display:block; font-weight:600; margin-bottom:.3rem; }
|
|
input, textarea, button { width:100%; border-radius:10px; border:1px solid #d8ccb7; padding:.65rem; font-size:.95rem; }
|
|
textarea { min-height: 110px; resize: vertical; }
|
|
button { background: var(--accent); color:white; border:none; font-weight:700; cursor:pointer; }
|
|
button:disabled { opacity:.6; cursor:default; }
|
|
.muted { color: var(--muted); font-size:.88rem; }
|
|
.out img { width:100%; border-radius:12px; border:1px solid #e4d8c3; background:white; }
|
|
@media (max-width: 760px) { .grid { grid-template-columns: 1fr; } }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class=\"wrap\">
|
|
<div class=\"card\">
|
|
<h1>FLUX.2 Klein Playground</h1>
|
|
<p class=\"muted\">Private single-user UI at <code>/</code>. Distilled defaults: 4 steps, guidance 1.0.</p>
|
|
<form id=\"form\" class=\"grid\">
|
|
<div class=\"full\">
|
|
<label for=\"prompt\">Prompt</label>
|
|
<textarea id=\"prompt\" required>A cinematic photo portrait, natural skin texture, realistic lighting, high detail.</textarea>
|
|
</div>
|
|
<div>
|
|
<label for=\"width\">Width (multiple of 16)</label>
|
|
<input id=\"width\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
|
|
</div>
|
|
<div>
|
|
<label for=\"height\">Height (multiple of 16)</label>
|
|
<input id=\"height\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
|
|
</div>
|
|
<div>
|
|
<label for=\"steps\">Inference steps</label>
|
|
<input id=\"steps\" type=\"number\" value=\"4\" min=\"1\" max=\"100\" />
|
|
</div>
|
|
<div>
|
|
<label for=\"guidance\">Guidance scale</label>
|
|
<input id=\"guidance\" type=\"number\" value=\"1.0\" step=\"0.1\" />
|
|
</div>
|
|
<div>
|
|
<label for=\"seed\">Seed</label>
|
|
<input id=\"seed\" type=\"number\" value=\"0\" min=\"0\" />
|
|
</div>
|
|
<div>
|
|
<label for=\"edit\">Edit image (optional)</label>
|
|
<input id=\"edit\" type=\"file\" accept=\"image/*\" />
|
|
</div>
|
|
<div class=\"full\">
|
|
<button id=\"submit\" type=\"submit\">Generate</button>
|
|
<p id=\"status\" class=\"muted\">Idle</p>
|
|
</div>
|
|
</form>
|
|
</div>
|
|
|
|
<div class=\"card out\" style=\"margin-top:1rem\">
|
|
<h1>Output</h1>
|
|
<img id=\"preview\" alt=\"Generated image preview\" />
|
|
<p class=\"muted\">Tip: right-click image to save.</p>
|
|
</div>
|
|
</div>
|
|
|
|
<script>
|
|
const form = document.getElementById('form');
|
|
const statusEl = document.getElementById('status');
|
|
const preview = document.getElementById('preview');
|
|
const button = document.getElementById('submit');
|
|
|
|
async function fileToDataUrl(file) {
|
|
return await new Promise((resolve, reject) => {
|
|
const reader = new FileReader();
|
|
reader.onload = () => resolve(reader.result);
|
|
reader.onerror = reject;
|
|
reader.readAsDataURL(file);
|
|
});
|
|
}
|
|
|
|
form.addEventListener('submit', async (event) => {
|
|
event.preventDefault();
|
|
button.disabled = true;
|
|
statusEl.textContent = 'Generating...';
|
|
try {
|
|
const prompt = document.getElementById('prompt').value;
|
|
const width = Number(document.getElementById('width').value);
|
|
const height = Number(document.getElementById('height').value);
|
|
const steps = Number(document.getElementById('steps').value);
|
|
const guidance = Number(document.getElementById('guidance').value);
|
|
const seed = Number(document.getElementById('seed').value);
|
|
const file = document.getElementById('edit').files[0];
|
|
const imageDataUrl = file ? await fileToDataUrl(file) : null;
|
|
|
|
const payload = {
|
|
model: 'black-forest-labs/FLUX.2-klein-9B',
|
|
prompt,
|
|
size: `${width}x${height}`,
|
|
num_inference_steps: steps,
|
|
guidance_scale: guidance,
|
|
seed,
|
|
response_format: 'b64_json',
|
|
imageDataUrl,
|
|
};
|
|
|
|
const resp = await fetch('/v1/images/generations', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify(payload),
|
|
});
|
|
|
|
if (!resp.ok) {
|
|
const text = await resp.text();
|
|
throw new Error(text || `HTTP ${resp.status}`);
|
|
}
|
|
|
|
const data = await resp.json();
|
|
const b64 = data?.data?.[0]?.b64_json;
|
|
if (!b64) throw new Error('No image returned');
|
|
preview.src = `data:image/png;base64,${b64}`;
|
|
statusEl.textContent = `Done. Seed: ${data.seed ?? seed}`;
|
|
} catch (error) {
|
|
statusEl.textContent = `Error: ${error.message}`;
|
|
} finally {
|
|
button.disabled = false;
|
|
}
|
|
});
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
|
|
@app.get("/health")
|
|
def health() -> dict[str, str]:
|
|
return {"status": "ok", "model": CHECKPOINT}
|
|
|
|
|
|
@app.post("/generate")
|
|
def generate(req: GenerateRequest) -> Response:
|
|
width, height = _normalize_dims(req.width, req.height)
|
|
seed = req.seed if not req.randomize_seed else random.randint(0, MAX_SEED)
|
|
image = _decode_image_data_url(req.image_data_url)
|
|
png_bytes, used_seed = _generate_png_bytes(
|
|
prompt=req.prompt,
|
|
width=width,
|
|
height=height,
|
|
num_inference_steps=req.num_inference_steps,
|
|
guidance_scale=req.guidance_scale,
|
|
seed=seed,
|
|
image=image,
|
|
)
|
|
return Response(
|
|
content=png_bytes,
|
|
media_type="image/png",
|
|
headers={"X-Seed": str(used_seed)},
|
|
)
|
|
|
|
|
|
@app.post("/v1/images/generations")
|
|
def compat_generate(req: CompatRequest) -> dict[str, Any]:
|
|
width, height = _parse_size(req.size, req.width, req.height)
|
|
seed_value = req.seed if req.seed is not None else 0
|
|
seed = seed_value if not req.randomize_seed else random.randint(0, MAX_SEED)
|
|
image = _decode_image_data_url(req.imageDataUrl)
|
|
|
|
png_bytes, used_seed = _generate_png_bytes(
|
|
prompt=req.prompt,
|
|
width=width,
|
|
height=height,
|
|
num_inference_steps=req.num_inference_steps or DEFAULT_STEPS,
|
|
guidance_scale=req.guidance_scale or DEFAULT_GUIDANCE,
|
|
seed=seed,
|
|
image=image,
|
|
)
|
|
b64_png = base64.b64encode(png_bytes).decode("ascii")
|
|
return {
|
|
"created": int(time.time()),
|
|
"model": CHECKPOINT,
|
|
"seed": used_seed,
|
|
"data": [{"b64_json": b64_png}],
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
parser = argparse.ArgumentParser(description="Private FLUX.2 Klein FastAPI server")
|
|
parser.add_argument("--model", choices=["4b", "9b"], default=MODEL_SIZE)
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=6006)
|
|
args = parser.parse_args()
|
|
|
|
MODEL_SIZE = args.model
|
|
CHECKPOINT = _checkpoint_for(MODEL_SIZE)
|
|
uvicorn.run(app, host=args.host, port=args.port)
|