123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import base64
|
|
import gradio as gr
|
|
import requests
|
|
from typing import List, Dict, Any, Tuple
|
|
|
|
from annotator.openpose import decode_json_as_poses, draw_poses
|
|
from scripts.controlnet_ui.modal import ModalInterface
|
|
from modules import shared
|
|
from scripts.logging import logger
|
|
|
|
|
|
def parse_data_url(data_url: str):
|
|
# Split the URL at the comma
|
|
media_type, data = data_url.split(",", 1)
|
|
|
|
# Check if the data is base64-encoded
|
|
assert ";base64" in media_type
|
|
|
|
# Decode the base64 data
|
|
return base64.b64decode(data)
|
|
|
|
|
|
def encode_data_url(json_string: str) -> str:
|
|
base64_encoded_json = base64.b64encode(json_string.encode("utf-8")).decode("utf-8")
|
|
return f"data:application/json;base64,{base64_encoded_json}"
|
|
|
|
|
|
class OpenposeEditor(object):
|
|
# Filename used when user click the download link.
|
|
download_file = "pose.json"
|
|
# URL the openpose editor is mounted on.
|
|
editor_url = "/openpose_editor_index"
|
|
|
|
def __init__(self) -> None:
|
|
self.render_button = None
|
|
self.download_link = None
|
|
self.modal = None
|
|
self.render()
|
|
|
|
def render(self):
|
|
# The hidden button to trigger a re-render of generated image.
|
|
self.render_button = gr.Button(visible=False, elem_classes=["cnet-render-pose"])
|
|
# The hidden element that stores the pose json for backend retrieval.
|
|
# The front-end javascript will write the edited JSON data to the element.
|
|
self.pose_input = gr.Textbox(visible=False, elem_classes=["cnet-pose-json"])
|
|
|
|
self.modal = ModalInterface(
|
|
# Use about:blank here as placeholder so that the iframe does not
|
|
# immediately navigate. Most of controlnet units do not need
|
|
# openpose editor active. Only navigate when the user first click
|
|
# 'Edit'. The navigation logic is in `openpose_editor.js`.
|
|
f'<iframe src="about:blank"></iframe>',
|
|
open_button_text="Edit",
|
|
open_button_classes=["cnet-edit-pose"],
|
|
open_button_extra_attrs=f'title="Send pose to {OpenposeEditor.editor_url} for edit."',
|
|
).create_modal(visible=False)
|
|
self.download_link = gr.HTML(
|
|
value="", visible=False, elem_classes=["cnet-download-pose"]
|
|
)
|
|
|
|
def register_callbacks(
|
|
self, generated_image: gr.Image, use_preview_as_input: gr.Checkbox
|
|
):
|
|
def render_pose(pose_url: str) -> Tuple[Dict, Dict]:
|
|
json_string = parse_data_url(pose_url)
|
|
poses, height, weight = decode_json_as_poses(json_string)
|
|
logger.info('Preview as input is enabled.')
|
|
return (
|
|
# Generated image.
|
|
gr.update(
|
|
value=draw_poses(
|
|
poses,
|
|
height,
|
|
weight,
|
|
draw_body=True,
|
|
draw_hand=True,
|
|
draw_face=True,
|
|
),
|
|
visible=True,
|
|
),
|
|
# Use preview as input.
|
|
gr.update(value=True),
|
|
)
|
|
|
|
self.render_button.click(
|
|
fn=render_pose,
|
|
inputs=[self.pose_input],
|
|
outputs=[generated_image, use_preview_as_input],
|
|
)
|
|
|
|
def outputs(self) -> List[Any]:
|
|
return [
|
|
self.download_link,
|
|
self.modal,
|
|
]
|
|
|
|
def update(self, json_string: str) -> List[Dict]:
|
|
"""
|
|
Called when there is a new JSON pose value generated by running
|
|
preprocessor.
|
|
|
|
Args:
|
|
json_string: The new JSON string generated by preprocessor.
|
|
|
|
Returns:
|
|
An gr.update event.
|
|
"""
|
|
hint = "Download the pose as .json file"
|
|
html = f"""<a href='{encode_data_url(json_string)}'
|
|
download='{OpenposeEditor.download_file}' title="{hint}">
|
|
JSON</a>"""
|
|
|
|
visible = json_string != ""
|
|
return [
|
|
# Download link update.
|
|
gr.update(value=html, visible=visible),
|
|
# Modal update.
|
|
gr.update(
|
|
visible=visible
|
|
and not shared.opts.data.get("controlnet_disable_openpose_edit", False)
|
|
),
|
|
]
|