user
49c0cf8587
- Add path traversal protection for model_name - Add file size limit (50MB) - Add TTL for results cache (5 min) with background cleanup - Fix bare except clause in client
237 lines
8.1 KiB
Python
237 lines
8.1 KiB
Python
import io
|
|
import os
|
|
import json
|
|
import torch
|
|
import uuid
|
|
import time
|
|
import threading
|
|
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
|
|
from fastapi.responses import Response, StreamingResponse
|
|
from PIL import Image
|
|
import numpy as np
|
|
from spandrel import ModelLoader
|
|
|
|
app = FastAPI(title="Remote Upscale Server", version="1.0.0")
|
|
|
|
MODELS_PATH = os.environ.get("UPSCALE_MODELS_PATH", "./models")
|
|
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB
|
|
RESULT_TTL = 300 # 5 minutes
|
|
|
|
loaded_models = {}
|
|
results_cache = {} # {result_id: {"data": bytes, "created": timestamp}}
|
|
|
|
|
|
def cleanup_old_results():
|
|
"""Remove results older than RESULT_TTL."""
|
|
now = time.time()
|
|
expired = [k for k, v in results_cache.items() if now - v["created"] > RESULT_TTL]
|
|
for k in expired:
|
|
del results_cache[k]
|
|
|
|
|
|
def start_cleanup_thread():
|
|
"""Background thread to clean up old results every 60 seconds."""
|
|
def cleanup_loop():
|
|
while True:
|
|
time.sleep(60)
|
|
cleanup_old_results()
|
|
thread = threading.Thread(target=cleanup_loop, daemon=True)
|
|
thread.start()
|
|
|
|
|
|
start_cleanup_thread()
|
|
|
|
|
|
def validate_model_name(model_name: str) -> str:
|
|
"""Validate model name to prevent path traversal attacks."""
|
|
# Check for path traversal attempts
|
|
if ".." in model_name or "/" in model_name or "\\" in model_name:
|
|
raise HTTPException(status_code=400, detail="Invalid model name")
|
|
|
|
model_path = os.path.join(MODELS_PATH, model_name)
|
|
|
|
# Verify the path is within MODELS_PATH
|
|
if not os.path.realpath(model_path).startswith(os.path.realpath(MODELS_PATH)):
|
|
raise HTTPException(status_code=400, detail="Invalid model name")
|
|
|
|
if not os.path.exists(model_path):
|
|
raise HTTPException(status_code=404, detail=f"Model not found: {model_name}")
|
|
|
|
return model_path
|
|
|
|
|
|
def get_model(model_name: str):
|
|
model_path = validate_model_name(model_name)
|
|
if model_name not in loaded_models:
|
|
print(f"Loading model: {model_path}")
|
|
model = ModelLoader().load_from_file(model_path).eval().cuda()
|
|
loaded_models[model_name] = model
|
|
return loaded_models[model_name]
|
|
|
|
|
|
@app.get("/models")
|
|
async def list_models():
|
|
"""List available upscale models."""
|
|
models = []
|
|
if os.path.exists(MODELS_PATH):
|
|
for f in os.listdir(MODELS_PATH):
|
|
if f.endswith(('.pth', '.pt', '.safetensors')):
|
|
models.append(f)
|
|
return {"models": sorted(models), "path": MODELS_PATH}
|
|
|
|
|
|
@app.post("/upscale")
|
|
async def upscale(
|
|
file: UploadFile = File(...),
|
|
model_name: str = Form(...),
|
|
tile_size: int = Form(512)
|
|
):
|
|
"""Upscale an image without progress streaming."""
|
|
image_data = await file.read()
|
|
if len(image_data) > MAX_FILE_SIZE:
|
|
raise HTTPException(status_code=413, detail=f"File too large. Max size: {MAX_FILE_SIZE // 1024 // 1024}MB")
|
|
image = Image.open(io.BytesIO(image_data)).convert("RGB")
|
|
|
|
img_np = np.array(image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).cuda()
|
|
|
|
model = get_model(model_name)
|
|
|
|
with torch.no_grad():
|
|
b, c, h, w = img_tensor.shape
|
|
if h * w < tile_size * tile_size * 4:
|
|
output = model(img_tensor)
|
|
else:
|
|
output = tiled_upscale(img_tensor, model, tile_size)
|
|
|
|
output_np = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
|
output_np = np.clip(output_np * 255, 0, 255).astype(np.uint8)
|
|
output_image = Image.fromarray(output_np)
|
|
|
|
img_bytes = io.BytesIO()
|
|
output_image.save(img_bytes, format="PNG")
|
|
img_bytes.seek(0)
|
|
|
|
return Response(content=img_bytes.read(), media_type="image/png")
|
|
|
|
|
|
@app.post("/upscale_stream")
|
|
async def upscale_stream(
|
|
file: UploadFile = File(...),
|
|
model_name: str = Form(...),
|
|
tile_size: int = Form(512)
|
|
):
|
|
"""Upscale an image with SSE progress streaming."""
|
|
image_data = await file.read()
|
|
if len(image_data) > MAX_FILE_SIZE:
|
|
raise HTTPException(status_code=413, detail=f"File too large. Max size: {MAX_FILE_SIZE // 1024 // 1024}MB")
|
|
image = Image.open(io.BytesIO(image_data)).convert("RGB")
|
|
|
|
img_np = np.array(image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).cuda()
|
|
|
|
model = get_model(model_name)
|
|
result_id = str(uuid.uuid4())
|
|
|
|
async def generate():
|
|
nonlocal img_tensor
|
|
with torch.no_grad():
|
|
b, c, h, w = img_tensor.shape
|
|
scale = model.scale
|
|
overlap = 32
|
|
|
|
tiles_y = max(1, (h + tile_size - overlap - 1) // (tile_size - overlap))
|
|
tiles_x = max(1, (w + tile_size - overlap - 1) // (tile_size - overlap))
|
|
total_tiles = tiles_y * tiles_x
|
|
|
|
if h * w < tile_size * tile_size * 4:
|
|
yield f"data: {json.dumps({'progress': 0, 'total': 1})}\n\n"
|
|
output = model(img_tensor)
|
|
yield f"data: {json.dumps({'progress': 1, 'total': 1})}\n\n"
|
|
else:
|
|
out_h, out_w = int(h * scale), int(w * scale)
|
|
output = torch.zeros((b, c, out_h, out_w), device=img_tensor.device)
|
|
count = torch.zeros((b, c, out_h, out_w), device=img_tensor.device)
|
|
|
|
current_tile = 0
|
|
for y in range(0, h, tile_size - overlap):
|
|
for x in range(0, w, tile_size - overlap):
|
|
y_end = min(y + tile_size, h)
|
|
x_end = min(x + tile_size, w)
|
|
|
|
tile = img_tensor[:, :, y:y_end, x:x_end]
|
|
upscaled_tile = model(tile)
|
|
|
|
out_y, out_x = int(y * scale), int(x * scale)
|
|
out_y_end, out_x_end = int(y_end * scale), int(x_end * scale)
|
|
|
|
output[:, :, out_y:out_y_end, out_x:out_x_end] += upscaled_tile
|
|
count[:, :, out_y:out_y_end, out_x:out_x_end] += 1
|
|
|
|
current_tile += 1
|
|
yield f"data: {json.dumps({'progress': current_tile, 'total': total_tiles})}\n\n"
|
|
|
|
output = output / count.clamp(min=1)
|
|
|
|
output_np = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
|
output_np = np.clip(output_np * 255, 0, 255).astype(np.uint8)
|
|
output_image = Image.fromarray(output_np)
|
|
|
|
img_bytes = io.BytesIO()
|
|
output_image.save(img_bytes, format="PNG")
|
|
results_cache[result_id] = {"data": img_bytes.getvalue(), "created": time.time()}
|
|
|
|
yield f"data: {json.dumps({'status': 'done', 'result_id': result_id})}\n\n"
|
|
|
|
return StreamingResponse(generate(), media_type="text/event-stream")
|
|
|
|
|
|
@app.get("/result/{result_id}")
|
|
async def get_result(result_id: str):
|
|
"""Fetch the result image after streaming completes."""
|
|
if result_id not in results_cache:
|
|
raise HTTPException(status_code=404, detail="Result not found or expired")
|
|
|
|
img_data = results_cache.pop(result_id)["data"]
|
|
return Response(content=img_data, media_type="image/png")
|
|
|
|
|
|
def tiled_upscale(img, model, tile_size, overlap=32):
|
|
"""Process image in tiles to handle large images."""
|
|
b, c, h, w = img.shape
|
|
scale = model.scale
|
|
|
|
out_h, out_w = int(h * scale), int(w * scale)
|
|
output = torch.zeros((b, c, out_h, out_w), device=img.device)
|
|
count = torch.zeros((b, c, out_h, out_w), device=img.device)
|
|
|
|
for y in range(0, h, tile_size - overlap):
|
|
for x in range(0, w, tile_size - overlap):
|
|
y_end = min(y + tile_size, h)
|
|
x_end = min(x + tile_size, w)
|
|
|
|
tile = img[:, :, y:y_end, x:x_end]
|
|
upscaled_tile = model(tile)
|
|
|
|
out_y, out_x = int(y * scale), int(x * scale)
|
|
out_y_end, out_x_end = int(y_end * scale), int(x_end * scale)
|
|
|
|
output[:, :, out_y:out_y_end, out_x:out_x_end] += upscaled_tile
|
|
count[:, :, out_y:out_y_end, out_x:out_x_end] += 1
|
|
|
|
return output / count.clamp(min=1)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Health check endpoint."""
|
|
return {
|
|
"status": "ok",
|
|
"device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8189)
|