Enhancement:support api use upscaler

pull/224/head
fengchengyang 2023-08-25 15:06:03 +08:00
parent 3176d477d7
commit d1b3a7e16e
1 changed files with 10 additions and 3 deletions

View File

@ -30,6 +30,12 @@ def get_full_model(model_name):
return model
return None
def get_upscaler(name):
for upscaler in shared.sd_upscalers:
if upscaler.name == name:
return upscaler
return None
def roop_api(_: gr.Blocks, app: FastAPI):
@app.post("/roop/image")
async def roop_image(
@ -37,6 +43,7 @@ def roop_api(_: gr.Blocks, app: FastAPI):
target_image: str = Body("",title="target image"),
face_index: list[int] = Body([0],title="face index"),
scale: int = Body(1,title="scale"),
upscaler: str = Body("None", title="Upscaler"),
upscale_visibility: float = Body(1,title="upscale visibility"),
face_restorer: str = Body("None",title="face restorer"),
restorer_visibility: float = Body(1,title="face restorer"),
@ -45,20 +52,20 @@ def roop_api(_: gr.Blocks, app: FastAPI):
s_image = api.decode_base64_to_image(source_image)
t_image = api.decode_base64_to_image(target_image)
f_index = set(face_index)
up_options = UpscaleOptions(scale=scale, upscale_visibility=upscale_visibility,face_restorer=get_face_restorer(face_restorer),restorer_visibility=restorer_visibility)
up_options = UpscaleOptions(scale=scale, upscaler=get_upscaler(upscaler), upscale_visibility=upscale_visibility,face_restorer=get_face_restorer(face_restorer),restorer_visibility=restorer_visibility)
use_model = get_full_model(model)
if use_model is None:
Exception("Model not found")
result = swap_face(s_image, t_image, use_model, f_index, up_options)
return {"image": api.encode_pil_to_base64(result.image())}
@app.get("/roop/models")
async def roop_models():
models = []
for model in get_models():
models.append(model.split("/")[-1])
return {"models": models}
try:
import modules.script_callbacks as script_callbacks