Split repo into apiserver/ and webserver/, add auth error handling and nginx proxy timeouts

2026-09-28 23:05:00 -05:00
parent 32149a88e1
commit 13dbe86f19
13 changed files with 419 additions and 149 deletions
@@ -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)