upload a cn

extensions-builtin/sd_forge_controlnet/scripts/api.py

@@ -0,0 +1,197 @@
+from typing import List, Optional
+
+import numpy as np
+from fastapi import FastAPI, Body
+from fastapi.exceptions import HTTPException
+from pydantic import BaseModel
+
+from PIL import Image
+
+import gradio as gr
+
+from modules.api.models import *
+from modules.api import api
+
+from scripts import external_code, global_state
+from scripts.processor import preprocessor_filters
+from scripts.logging import logger
+from annotator.openpose import draw_poses, decode_json_as_poses
+from annotator.openpose.animalpose import draw_animalposes
+
+
+def encode_to_base64(image):
+    if isinstance(image, str):
+        return 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:
+        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/version")
+    async def version():
+        return {"version": external_code.get_api_version()}
+
+    @app.get("/controlnet/model_list")
+    async def model_list(update: bool = True):
+        up_to_date_model_list = external_code.get_models(update=update)
+        logger.debug(up_to_date_model_list)
+        return {"model_list": up_to_date_model_list}
+
+    @app.get("/controlnet/module_list")
+    async def module_list(alias_names: bool = False):
+        _module_list = external_code.get_modules(alias_names)
+        logger.debug(_module_list)
+
+        return {
+            "module_list": _module_list,
+            "module_detail": external_code.get_modules_detail(alias_names),
+        }
+
+    @app.get("/controlnet/control_types")
+    async def control_types():
+        def format_control_type(
+            filtered_preprocessor_list,
+            filtered_model_list,
+            default_option,
+            default_model,
+        ):
+            return {
+                "module_list": filtered_preprocessor_list,
+                "model_list": filtered_model_list,
+                "default_option": default_option,
+                "default_model": default_model,
+            }
+
+        return {
+            "control_types": {
+                control_type: format_control_type(
+                    *global_state.select_control_type(control_type)
+                )
+                for control_type in preprocessor_filters.keys()
+            }
+        }
+
+    @app.get("/controlnet/settings")
+    async def settings():
+        max_models_num = external_code.get_max_models_num()
+        return {"control_net_unit_count": max_models_num}
+
+    cached_cn_preprocessors = global_state.cache_preprocessors(
+        global_state.cn_preprocessor_modules
+    )
+
+    @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"),
+        low_vram: bool = Body(False, title="Low vram"),
+    ):
+        controlnet_module = global_state.reverse_preprocessor_aliases.get(
+            controlnet_module, controlnet_module
+        )
+
+        if controlnet_module not in cached_cn_preprocessors:
+            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.info(
+            f"Detecting {str(len(controlnet_input_images))} images with the {controlnet_module} module."
+        )
+
+        results = []
+        poses = []
+
+        processor_module = cached_cn_preprocessors[controlnet_module]
+
+        for input_image in controlnet_input_images:
+            img = external_code.to_base64_nparray(input_image)
+
+            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,
+                    res=controlnet_processor_res,
+                    thr_a=controlnet_threshold_a,
+                    thr_b=controlnet_threshold_b,
+                    json_pose_callback=json_acceptor.accept,
+                    low_vram=low_vram,
+                )[0]
+            )
+
+            if "openpose" in controlnet_module:
+                assert json_acceptor.value is not None
+                poses.append(json_acceptor.value)
+
+        global_state.cn_preprocessor_unloadable.get(controlnet_module, lambda: None)()
+        results64 = list(map(encode_to_base64, results))
+        res = {"images": results64, "info": "Success"}
+        if poses:
+            res["poses"] = poses
+
+        return res
+
+    class Person(BaseModel):
+        pose_keypoints_2d: List[float]
+        hand_right_keypoints_2d: Optional[List[float]]
+        hand_left_keypoints_2d: Optional[List[float]]
+        face_keypoints_2d: Optional[List[float]]
+
+    class PoseData(BaseModel):
+        people: List[Person]
+        canvas_width: int
+        canvas_height: int
+
+    @app.post("/controlnet/render_openpose_json")
+    async def render_openpose_json(
+        pose_data: List[PoseData] = Body([], title="Pose json files to render.")
+    ):
+        if not pose_data:
+            return {"info": "No pose data detected."}
+        else:
+
+            def draw(poses, animals, H, W):
+                if poses:
+                    assert len(animals) == 0
+                    return draw_poses(poses, H, W)
+                else:
+                    return draw_animalposes(animals, H, W)
+
+            return {
+                "images": [
+                    encode_to_base64(draw(*decode_json_as_poses(pose.dict())))
+                    for pose in pose_data
+                ],
+                "info": "Success",
+            }
+
+
+try:
+    import modules.script_callbacks as script_callbacks
+
+    script_callbacks.on_app_started(controlnet_api)
+except:
+    pass