Skip to content

Commit

Permalink
UPDATE: API RGBA support, GET facemodels show
Browse files Browse the repository at this point in the history
FR: #394
D: #384
  • Loading branch information
Gourieff committed Mar 12, 2024
1 parent 0185d7a commit 9989610
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 4 deletions.
2 changes: 1 addition & 1 deletion example/api_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
print(im)

img_bytes = io.BytesIO()
im.save(img_bytes, format='PNG')
im.save(img_bytes, format='PNG')
img_base64 = base64.b64encode(img_bytes.getvalue()).decode('utf-8')

# ReActor arguments:
Expand Down
43 changes: 40 additions & 3 deletions scripts/reactor_api.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,18 @@
'''
Thanks SpenserCai for the original version of the roop api script
-----------------------------------
--- ReActor External API v1.0.6 ---
--- ReActor External API v1.0.7 ---
-----------------------------------
'''
import os, glob
from datetime import datetime, date
from fastapi import FastAPI, Body
# from fastapi.exceptions import HTTPException
# from io import BytesIO
# from PIL import Image
# import base64
# import numpy as np
# import cv2

# from modules.api.models import *
from modules import scripts, shared
Expand All @@ -16,6 +22,7 @@

from scripts.reactor_swapper import EnhancementOptions, swap_face, DetectionOptions
from scripts.reactor_logger import logger
from scripts.reactor_helpers import get_facemodels

# XYZ init:
from scripts.reactor_xyz import run
Expand Down Expand Up @@ -61,6 +68,19 @@ def get_full_model(model_name):
return model
return None

# def decode_base64_to_image_rgba(encoding):
# if encoding.startswith("data:image/"):
# encoding = encoding.split(";")[1].split(",")[1]
# try:
# im_bytes = base64.b64decode(encoding)
# im_arr = np.frombuffer(im_bytes, dtype=np.uint8) # im_arr is one-dim Numpy array
# img = cv2.imdecode(im_arr, flags=cv2.IMREAD_UNCHANGED)
# img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
# image = Image.fromarray(img, mode="RGBA")
# return image
# except Exception as e:
# raise HTTPException(status_code=500, detail="Invalid encoded image") from e

def reactor_api(_: gr.Blocks, app: FastAPI):
@app.post("/reactor/image")
async def reactor_image(
Expand Down Expand Up @@ -92,6 +112,12 @@ async def reactor_image(
):
s_image = api.decode_base64_to_image(source_image) if select_source == 0 else None
t_image = api.decode_base64_to_image(target_image)

if t_image.mode == 'RGBA':
_, _, _, alpha = t_image.split()
else:
alpha = None

sf_index = source_faces_index
f_index = face_index
gender_s = gender_source
Expand All @@ -106,15 +132,21 @@ async def reactor_image(
if use_model is None:
Exception("Model not found")
result = swap_face(s_image, t_image, use_model, sf_index, f_index, up_options, gender_s, gender_t, True, True, device, mask_face, select_source, face_model, source_folder, None, random_image,det_options)
result_img = result[0]

if alpha is not None:
result_img = result_img.convert("RGBA")
result_img.putalpha(alpha)

if save_to_file == 1:
if result_file_path == "":
result_file_path = default_file_path()
try:
result[0].save(result_file_path, format='PNG')
result_img.save(result_file_path, format='PNG')
logger.status("Result has been saved to: %s", result_file_path)
except Exception as e:
logger.error("Error while saving result: %s",e)
return {"image": api.encode_pil_to_base64(result[0])}
return {"image": api.encode_pil_to_base64(result_img)}

@app.get("/reactor/models")
async def reactor_models():
Expand All @@ -125,6 +157,11 @@ async def reactor_models():
async def reactor_upscalers():
names = [upscaler.name for upscaler in shared.sd_upscalers]
return {"upscalers": names}

@app.get("/reactor/facemodels")
async def reactor_facemodels():
facemodels = [os.path.split(model)[1].split(".")[0] for model in get_facemodels()]
return {"facemodels": facemodels}

try:
import modules.script_callbacks as script_callbacks
Expand Down

0 comments on commit 9989610

Please sign in to comment.