Split repo into apiserver/ and webserver/, add auth error handling and nginx proxy timeouts
This commit is contained in:
@@ -0,0 +1,26 @@
|
||||
# Uses the exact same verified environment as your Dev Container
|
||||
FROM nvcr.io/nvidia/pytorch:25.02-py3
|
||||
|
||||
# Install API and Flux-specific dependencies
|
||||
# Note: We skip the torch versions in your list because the NVIDIA base
|
||||
# already includes the optimized sm_120 version we need.
|
||||
RUN pip3 install --no-cache-dir \
|
||||
fastapi \
|
||||
uvicorn \
|
||||
pydantic \
|
||||
transformers==4.46.1 \
|
||||
diffusers==0.35.2 \
|
||||
accelerate==0.30.1 \
|
||||
sentencepiece \
|
||||
protobuf \
|
||||
bcrypt
|
||||
|
||||
# Set up the application
|
||||
WORKDIR /app
|
||||
COPY api.py /app/api.py
|
||||
|
||||
# Ensure GPU visibility (matching your devcontainer.json)
|
||||
ENV NVIDIA_VISIBLE_DEVICES=all
|
||||
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
|
||||
|
||||
CMD ["python3", "api.py"]
|
||||
@@ -0,0 +1,191 @@
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import time
|
||||
import hmac
|
||||
import hashlib
|
||||
import secrets
|
||||
from io import BytesIO
|
||||
from fastapi import FastAPI, Request, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, model_validator
|
||||
from typing import Any
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
import bcrypt
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
|
||||
hf_token = os.environ.get("HF_TOKEN")
|
||||
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-schnell",
|
||||
torch_dtype=torch.bfloat16,
|
||||
token=hf_token
|
||||
)
|
||||
|
||||
# Offload execution stages sequentially to maintain 16GB headroom
|
||||
pipe.enable_sequential_cpu_offload()
|
||||
pipe.vae.enable_tiling()
|
||||
|
||||
# Auth config
|
||||
USERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'users.json')
|
||||
SESSION_SECRET = os.environ.get('SESSION_SECRET', secrets.token_hex(32))
|
||||
SESSION_DURATION = 86400 # 24 hours
|
||||
|
||||
|
||||
def load_users():
|
||||
with open(USERS_FILE) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def create_session(username: str) -> str:
|
||||
payload = f"{username}:{int(time.time())}"
|
||||
signature = hmac.new(
|
||||
SESSION_SECRET.encode(),
|
||||
payload.encode(),
|
||||
hashlib.sha256
|
||||
).hexdigest()
|
||||
return f"{payload}.{signature}"
|
||||
|
||||
|
||||
def verify_session(session: str) -> str | None:
|
||||
try:
|
||||
payload, sig = session.rsplit('.', 1)
|
||||
expected_sig = hmac.new(
|
||||
SESSION_SECRET.encode(),
|
||||
payload.encode(),
|
||||
hashlib.sha256
|
||||
).hexdigest()
|
||||
if not hmac.compare_digest(sig, expected_sig):
|
||||
return None
|
||||
username, timestamp = payload.split(':', 1)
|
||||
if time.time() - int(timestamp) > SESSION_DURATION:
|
||||
return None
|
||||
users = load_users()
|
||||
if username not in users:
|
||||
return None
|
||||
return username
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class ImageRequest(BaseModel):
|
||||
prompt: str
|
||||
model: str = "flux-schnell"
|
||||
width: int = 1024
|
||||
height: int = 1024
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_dimensions(cls, data: Any) -> Any:
|
||||
if isinstance(data, dict):
|
||||
width = data.get("width", 1024)
|
||||
height = data.get("height", 1024)
|
||||
if width % 32 != 0 or height % 32 != 0:
|
||||
raise ValueError("Width and height must be multiples of 32")
|
||||
if width > 2048 or height > 2048:
|
||||
raise ValueError("Maximum width and height is 2048")
|
||||
return data
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def auth_middleware(request: Request, call_next):
|
||||
# Allow auth endpoints and static files through
|
||||
if request.url.path.startswith(("/auth/login", "/auth/logout", "/auth/check")):
|
||||
return await call_next(request)
|
||||
|
||||
# Check session cookie for API requests
|
||||
session = request.cookies.get("session")
|
||||
username = verify_session(session) if session else None
|
||||
|
||||
if not username:
|
||||
return JSONResponse(
|
||||
status_code=401,
|
||||
content={"detail": "Not authenticated"}
|
||||
)
|
||||
|
||||
request.state.username = username
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@app.post("/auth/login")
|
||||
async def login(request: Request):
|
||||
data = await request.json()
|
||||
username = data.get('username', '')
|
||||
password = data.get('password', '')
|
||||
users = load_users()
|
||||
|
||||
if username not in users:
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
|
||||
if not bcrypt.checkpw(password.encode(), users[username].encode()):
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
|
||||
session = create_session(username)
|
||||
response = JSONResponse(content={"username": username})
|
||||
response.set_cookie(
|
||||
key="session",
|
||||
value=session,
|
||||
httponly=True,
|
||||
samesite="strict",
|
||||
max_age=SESSION_DURATION,
|
||||
path="/",
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
@app.post("/auth/logout")
|
||||
async def logout():
|
||||
response = JSONResponse(content={"message": "Logged out"})
|
||||
response.set_cookie("session", "", expires=0, path="/")
|
||||
return response
|
||||
|
||||
|
||||
@app.get("/auth/check")
|
||||
async def check(request: Request, session: str | None = None):
|
||||
username = verify_session(session) if session else None
|
||||
if not username:
|
||||
raise HTTPException(401, "Not authenticated")
|
||||
return {"username": username}
|
||||
|
||||
|
||||
@app.post("/v1/images/generations")
|
||||
async def generate(request: Request):
|
||||
try:
|
||||
body = await request.json()
|
||||
req = ImageRequest(**body)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
try:
|
||||
result = pipe(
|
||||
req.prompt,
|
||||
guidance_scale=0.0,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=256,
|
||||
width=req.width,
|
||||
height=req.height,
|
||||
).images[0]
|
||||
|
||||
buffered = BytesIO()
|
||||
result.save(buffered, format="PNG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode()
|
||||
|
||||
return {"data": [{"b64_json": img_str}]}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=5005)
|
||||
@@ -0,0 +1,27 @@
|
||||
services:
|
||||
flux-api:
|
||||
image: flux-image:1.0
|
||||
container_name: flux-api
|
||||
ports:
|
||||
- "5005:5005"
|
||||
ipc: host
|
||||
ulimits:
|
||||
memlock: -1
|
||||
stack: 67108864
|
||||
environment:
|
||||
- NVIDIA_VISIBLE_DEVICES=all
|
||||
- NVIDIA_DRIVER_CAPABILITIES=compute,utility
|
||||
- HF_TOKEN=${HF_TOKEN}
|
||||
- HF_HOME=/root/.cache/huggingface
|
||||
- SESSION_SECRET=${SESSION_SECRET}
|
||||
volumes:
|
||||
- /opt/huggingface_cache:/root/.cache/huggingface
|
||||
- /home/sysuser/Development/flux-api/users.json:/app/users.json
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
restart: unless-stopped
|
||||
Reference in New Issue
Block a user