123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122 |
- 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)
- ),
- ]
|