chore: Update everything
This commit is contained in:
@@ -0,0 +1,461 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user