|
from typing import List |
|
|
|
import numpy as np |
|
from fastapi import FastAPI, Body |
|
from fastapi.exceptions import HTTPException |
|
from PIL import Image |
|
import gradio as gr |
|
|
|
from modules.api import api |
|
from .global_state import ( |
|
get_all_preprocessor_names, |
|
get_all_controlnet_names, |
|
get_preprocessor, |
|
) |
|
from .utils import judge_image_type |
|
from .logging import logger |
|
|
|
|
|
def encode_to_base64(image): |
|
if isinstance(image, str): |
|
return image |
|
elif not judge_image_type(image): |
|
return "Detect result is not image" |
|
elif isinstance(image, Image.Image): |
|
return api.encode_pil_to_base64(image) |
|
elif isinstance(image, np.ndarray): |
|
return encode_np_to_base64(image) |
|
else: |
|
logger.warn("Unable to encode image.") |
|
return "" |
|
|
|
|
|
def encode_np_to_base64(image): |
|
pil = Image.fromarray(image) |
|
return api.encode_pil_to_base64(pil) |
|
|
|
|
|
def controlnet_api(_: gr.Blocks, app: FastAPI): |
|
@app.get("/controlnet/model_list") |
|
async def model_list(): |
|
up_to_date_model_list = get_all_controlnet_names() |
|
logger.debug(up_to_date_model_list) |
|
return {"model_list": up_to_date_model_list} |
|
|
|
@app.get("/controlnet/module_list") |
|
async def module_list(): |
|
module_list = get_all_preprocessor_names() |
|
logger.debug(module_list) |
|
|
|
return { |
|
"module_list": module_list, |
|
|
|
|
|
} |
|
|
|
@app.post("/controlnet/detect") |
|
async def detect( |
|
controlnet_module: str = Body("none", title="Controlnet Module"), |
|
controlnet_input_images: List[str] = Body([], title="Controlnet Input Images"), |
|
controlnet_processor_res: int = Body( |
|
512, title="Controlnet Processor Resolution" |
|
), |
|
controlnet_threshold_a: float = Body(64, title="Controlnet Threshold a"), |
|
controlnet_threshold_b: float = Body(64, title="Controlnet Threshold b"), |
|
): |
|
processor_module = get_preprocessor(controlnet_module) |
|
if processor_module is None: |
|
raise HTTPException(status_code=422, detail="Module not available") |
|
|
|
if len(controlnet_input_images) == 0: |
|
raise HTTPException(status_code=422, detail="No image selected") |
|
|
|
logger.debug( |
|
f"Detecting {str(len(controlnet_input_images))} images with the {controlnet_module} module." |
|
) |
|
|
|
results = [] |
|
poses = [] |
|
|
|
for input_image in controlnet_input_images: |
|
img = np.array(api.decode_base64_to_image(input_image)).astype('uint8') |
|
|
|
class JsonAcceptor: |
|
def __init__(self) -> None: |
|
self.value = None |
|
|
|
def accept(self, json_dict: dict) -> None: |
|
self.value = json_dict |
|
|
|
json_acceptor = JsonAcceptor() |
|
|
|
results.append( |
|
processor_module( |
|
img, |
|
resolution=controlnet_processor_res, |
|
slider_1=controlnet_threshold_a, |
|
slider_2=controlnet_threshold_b, |
|
json_pose_callback=json_acceptor.accept, |
|
) |
|
) |
|
|
|
if "openpose" in controlnet_module: |
|
assert json_acceptor.value is not None |
|
poses.append(json_acceptor.value) |
|
|
|
results64 = [encode_to_base64(img) for img in results] |
|
res = {"images": results64, "info": "Success"} |
|
if poses: |
|
res["poses"] = poses |
|
|
|
return res |
|
|
|
|
