Enhancement:support api use upscaler
parent
3176d477d7
commit
d1b3a7e16e
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue