87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
from fastapi import FastAPI, BackgroundTasks, HTTPException, Form
|
|
from fastapi.responses import FileResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
import subprocess
|
|
import os
|
|
from uuid import uuid4
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI()
|
|
|
|
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "output")
|
|
SD_PATH = os.environ.get("SD_PATH", "./sd")
|
|
MODEL_DIR = os.environ.get("MODEL_DIR", "./models")
|
|
MODEL_NAME = os.environ.get(
|
|
"MODEL_NAME", "v1-5-pruned-emaonly-ggml-model-q5_0.bin")
|
|
BASE_URL = os.environ.get("BASE_URL", "http://localhost:8000")
|
|
|
|
|
|
class Payload(BaseModel):
|
|
prompt: str
|
|
neg_prompt: str
|
|
seed: int
|
|
steps: int
|
|
width: int
|
|
height: int
|
|
|
|
|
|
# Create the OUTPUT_DIR directory if it does not exist
|
|
if not os.path.exists(OUTPUT_DIR):
|
|
os.makedirs(OUTPUT_DIR)
|
|
|
|
# Create the OUTPUT_DIR directory if it does not exist
|
|
if not os.path.exists(MODEL_DIR):
|
|
os.makedirs(MODEL_DIR)
|
|
|
|
# Serve files from the "files" directory
|
|
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
|
|
|
|
|
|
def run_command(payload: Payload, filename: str):
|
|
# Construct the command based on your provided example
|
|
command = [SD_PATH,
|
|
"--model", f'{os.path.join(MODEL_DIR, MODEL_NAME)}',
|
|
"--prompt", f'"{payload.prompt}"',
|
|
"--negative-prompt", f'"{payload.neg_prompt}"',
|
|
"--height", str(payload.height),
|
|
"--width", str(payload.width),
|
|
"--steps", str(payload.steps),
|
|
"--seed", str(payload.seed),
|
|
"--mode", 'txt2img',
|
|
"-o", f'{os.path.join(OUTPUT_DIR, filename)}',
|
|
]
|
|
|
|
try:
|
|
subprocess.run(command)
|
|
except subprocess.CalledProcessError:
|
|
raise HTTPException(
|
|
status_code=500, detail="Failed to execute the command.")
|
|
|
|
|
|
@app.post("/inferences/txt2img")
|
|
async def run_inference(background_tasks: BackgroundTasks, payload: Payload):
|
|
# Generate a unique filename using uuid4()
|
|
filename = f"{uuid4()}.png"
|
|
|
|
# We will use background task to run the command so it won't block
|
|
# background_tasks.add_task(run_command, payload, filename)
|
|
run_command(payload, filename)
|
|
|
|
# Return the expected path of the output file
|
|
return {"url": f'{BASE_URL}/serve/{filename}'}
|
|
|
|
|
|
@app.get("/serve/{filename}")
|
|
async def serve_file(filename: str):
|
|
file_path = os.path.join(OUTPUT_DIR, filename)
|
|
|
|
if os.path.exists(file_path):
|
|
return FileResponse(file_path)
|
|
else:
|
|
raise HTTPException(status_code=404, detail="File not found")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|