feature: api server (alpha)
`aimg run-api-server` Proof of concept for nowpull/335/head
Bryce
1 year ago
committed by
Bryce Drennan
parent
38ac0b7f54
commit
14739bc90b
@ -0,0 +1,14 @@
|
||||
import click
|
||||
|
||||
|
||||
@click.command("run-api-server")
|
||||
def run_api_server_cmd():
|
||||
"""Run a HTTP API server."""
|
||||
import uvicorn
|
||||
|
||||
from imaginairy.http.app import app
|
||||
from imaginairy.log_utils import configure_logging
|
||||
|
||||
configure_logging()
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
|
||||
@ -0,0 +1,81 @@
|
||||
from asyncio import Lock
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, Query
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from imaginairy import ImaginePrompt, imagine
|
||||
from imaginairy.log_utils import configure_logging
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
lock = Lock()
|
||||
|
||||
|
||||
class ImagineWebPrompt(BaseModel):
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
prompt: Optional[str]
|
||||
negative_prompt: Optional[str]
|
||||
prompt_strength: float = 7.5
|
||||
# init_image: Optional[Union[LazyLoadingImage, str]]
|
||||
init_image_strength: Optional[float] = None
|
||||
# control_inputs: Optional[List[ControlInput]] = None
|
||||
mask_prompt: Optional[str] = None
|
||||
# mask_image: Optional[Union[LazyLoadingImage, str]] = None
|
||||
mask_mode: str = "replace"
|
||||
mask_modify_original: bool = True
|
||||
outpaint: Optional[str] = None
|
||||
seed: Optional[int] = None
|
||||
steps: Optional[int] = None
|
||||
height: Optional[int] = None
|
||||
width: Optional[int] = None
|
||||
upscale: bool = False
|
||||
fix_faces: bool = False
|
||||
fix_faces_fidelity: float = 0.2
|
||||
# sampler_type: str = Field(..., alias='config.DEFAULT_SAMPLER') # update the alias based on actual config field name
|
||||
conditioning: Optional[str] = None
|
||||
tile_mode: str = ""
|
||||
allow_compose_phase: bool = True
|
||||
# model: str = Field(..., alias='config.DEFAULT_MODEL') # update the alias based on actual config field name
|
||||
model_config_path: Optional[str] = None
|
||||
is_intermediate: bool = False
|
||||
collect_progress_latents: bool = False
|
||||
caption_text: str = ""
|
||||
|
||||
|
||||
def generate_image(prompt: ImagineWebPrompt):
|
||||
prompt = ImaginePrompt(prompt.prompt)
|
||||
result = next(imagine([prompt]))
|
||||
return result.images["generated"]
|
||||
|
||||
|
||||
@app.post("/api/imagine")
|
||||
async def imagine_endpoint(prompt: ImagineWebPrompt):
|
||||
async with lock:
|
||||
img = await run_in_threadpool(generate_image, prompt)
|
||||
img_io = BytesIO()
|
||||
img.save(img_io, "JPEG")
|
||||
img_io.seek(0)
|
||||
return StreamingResponse(img_io, media_type="image/jpg")
|
||||
|
||||
|
||||
@app.get("/api/imagine")
|
||||
async def imagine_get_endpoint(text: str = Query(...)):
|
||||
async with lock:
|
||||
img = await run_in_threadpool(generate_image, ImagineWebPrompt(prompt=text))
|
||||
img_io = BytesIO()
|
||||
img.save(img_io, "JPEG")
|
||||
img_io.seek(0)
|
||||
return StreamingResponse(img_io, media_type="image/jpg")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
configure_logging()
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
|
||||
Loading…
Reference in New Issue