diff --git a/scripts/api.py b/scripts/api.py index 3943fa9..590db18 100644 --- a/scripts/api.py +++ b/scripts/api.py @@ -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