| | import os |
| | from io import BytesIO |
| |
|
| | import gradio as gr |
| | import grpc |
| | from PIL import Image |
| | import pandas as pd |
| |
|
| | from inference_pb2 import SFERequest, SFEResponse, SFERequestMask, SFEResponseMask |
| | from inference_pb2_grpc import SFEServiceStub |
| |
|
| |
|
| | PREDEFINED_EDITINGS_DATA = { |
| | "glasses": ([-20.0, 30.0], False), |
| | "smile": ([-10.0, 10.0], False), |
| | "makeup": ([-10.0, 15.0], False), |
| | "eye_openness": ([-45.0, 30.0], True), |
| | "trimmed_beard": ([-30.0, 30.0], True), |
| | "face_roundness": ([-20.0, 15.0], False), |
| | "nose_length": ([-30.0, 30.0], True), |
| | "eyebrow_thickness": ([-20.0, 20.0], True), |
| | "displeased": ([-10.0, 10.0], False), |
| | "age": ([-10.0, 10.0], False), |
| | "rotation": ([-7.0, 7.0], False), |
| | "afro": ([0, 0.14], False), |
| | "angry": ([0, 0.14], False), |
| | "bobcut": ([0, 0.18], False), |
| | "bowlcut": ([0, 0.14], False), |
| | "mohawk": ([0, 0.1], False), |
| | "curly_hair": ([0, 0.12], False), |
| | "purple_hair": ([0, 0.12], False), |
| | "surprised": ([0, 0.1], False), |
| | "beyonce": ([0, 0.12], False), |
| | "hilary_clinton": ([0, 0.1], False), |
| | "depp": ([0, 0.12], False), |
| | "taylor_swift": ([0, 0.1], False), |
| | "trump": ([0, 0.1], False), |
| | "zuckerberg": ([0, 0.1], False), |
| | "black hair": ([-7.0, 10.0], False), |
| | "blond hair": ([-7.0, 10.0], True), |
| | "grey hair": ([-7.0, 7.0], True), |
| | "wavy hair": ([-7.0, 7.0], False), |
| | "receding hairline": ([-10.0, 10.0], True), |
| | "sideburns": ([-7.0, 7.0], True), |
| | "goatee": ([-7.0, 7.0], True), |
| | "gender swap": ([-10.0, 7.0], False) |
| | } |
| |
|
| | DIRECTIONS_NAME_SWAP = { |
| | "smile" : "fs_smiling", |
| | "glasses": "fs_glasses", |
| | "makeup": "fs_makeup", |
| | "gender swap": "gender" |
| | } |
| |
|
| |
|
| | def denormalize_power(direction_name, directon_power): |
| | if direction_name not in PREDEFINED_EDITINGS_DATA: |
| | return directon_power |
| | original_range, is_reversed = PREDEFINED_EDITINGS_DATA[direction_name] |
| | if directon_power > 0: |
| | normalized = directon_power / 15 * abs(original_range[1]) |
| | else: |
| | normalized = directon_power / 15 * abs(original_range[0]) |
| |
|
| | if is_reversed: |
| | normalized = -normalized |
| | return normalized |
| |
|
| |
|
| | def get_bytes(img): |
| | if img is None: |
| | return img |
| |
|
| | buffered = BytesIO() |
| | img.save(buffered, format="JPEG") |
| | return buffered.getvalue() |
| |
|
| |
|
| | def bytes_to_image(image: bytes) -> Image.Image: |
| | image = Image.open(BytesIO(image)) |
| | return image |
| |
|
| |
|
| | def edit_image(orig_image, edit_direction, edit_power, align, mask, progress=gr.Progress(track_tqdm=True)): |
| | if edit_direction in DIRECTIONS_NAME_SWAP: |
| | edit_direction = DIRECTIONS_NAME_SWAP[edit_direction] |
| | if not orig_image: |
| | return gr.update(visible=False), gr.update(visible=False), gr.update(value="Need to upload an input image ❗", visible=True), gr.update(visible=False), gr.update(visible=False) |
| |
|
| | orig_image_bytes = get_bytes(orig_image) |
| | mask_bytes = get_bytes(mask) |
| | if mask_bytes is None: |
| | mask_bytes = b"mask" |
| |
|
| | edit_power = denormalize_power(edit_direction, edit_power) |
| |
|
| | with grpc.insecure_channel(os.environ["SERVER"]) as channel: |
| | stub = SFEServiceStub(channel) |
| |
|
| | output: SFEResponse = stub.edit( |
| | SFERequest(orig_image=orig_image_bytes, direction=edit_direction, power=edit_power, align=align, mask=mask_bytes, use_cache=True) |
| | ) |
| |
|
| | if output.image == b"aligner error": |
| | return gr.update(visible=False), gr.update(visible=False), gr.update(value="Face aligner can not find face in your image 😢 Try to upload another one", visible=True), gr.update(visible=False), gr.update(visible=False), |
| |
|
| | output_edited = bytes_to_image(output.image) |
| | output_inv = bytes_to_image(output.inv_image) |
| | if not align: |
| | return gr.update(value=output_edited, visible=True), gr.update(value=output_inv, visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) |
| |
|
| | output_aligned = bytes_to_image(output.aligned) |
| | output_unaligned = bytes_to_image(output.unaligned) |
| | return gr.update(value=output_edited, visible=True), gr.update(value=output_inv, visible=True), gr.update(visible=False), gr.update(value=output_aligned, visible=True), gr.update(value=output_unaligned, visible=True) |
| |
|
| |
|
| | def edit_image_clip(orig_image, neutral_prompt, target_prompt, disentanglement, edit_power, align, mask, edit_method, progress=gr.Progress(track_tqdm=True)): |
| | if edit_method == "StyleClip": |
| | edit_direction = "_".join(["styleclip_global", neutral_prompt, target_prompt, str(disentanglement)]) |
| | else: |
| | edit_power = edit_power / 10 |
| | disentanglement = disentanglement / 3 |
| | edit_direction = "_".join(["deltaedit", neutral_prompt, target_prompt, str(disentanglement)]) |
| | return edit_image(orig_image, edit_direction, edit_power, align, mask, progress=None) |
| |
|
| |
|
| | def get_mask(input_image, align, mask_trashhold, progress=gr.Progress(track_tqdm=True)): |
| | if not input_image: |
| | return gr.update(visible=False), gr.update(value="Need to upload an input image ❗", visible=True) |
| |
|
| | input_image_bytes = get_bytes(input_image) |
| |
|
| | with grpc.insecure_channel(os.environ["SERVER"]) as channel: |
| | stub = SFEServiceStub(channel) |
| |
|
| | output: SFEResponseMask = stub.generate_mask( |
| | SFERequestMask(orig_image=input_image_bytes, trashold=mask_trashhold, align=align, use_cache=True) |
| | ) |
| | if output.mask == b"aligner error": |
| | return gr.update(visible=False), gr.update(value="Face aligner can not find face in your image 😢 Try to upload another one", visible=True) |
| |
|
| | if output.mask == b"masker face parser error": |
| | return gr.update(visible=False), gr.update(value="Masker's face detector can't find face in your image 😢 Try to upload another one", visible=True) |
| |
|
| | output_mask = bytes_to_image(output.mask) |
| | return gr.update(value=output_mask, visible=True), gr.update(visible=False) |
| |
|
| |
|
| | def get_demo(): |
| | editings_table = pd.read_csv("editings_table.csv") |
| | editings_table = editings_table.style.set_properties(**{"text-align": "center"}) |
| | editings_table = editings_table.set_table_styles([dict(selector="th", props=[("text-align", "center")])]) |
| |
|
| | with gr.Blocks() as demo: |
| | gr.Markdown("## StyleFeatureEditor") |
| | gr.Markdown( |
| | '<div style="display: flex; align-items: center; gap: 10px;">' |
| | '<span>Official Gradio demo for StyleFeatureEditor:</span>' |
| | '<a href="https://arxiv.org/abs/2406.10601"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>' |
| | '<a href="https://github.com/AIRI-Institute/StyleFeatureEditor"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>' |
| | '<a href="https://huggingface.co/AIRI-Institute/StyleFeatureEditor"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>' |
| | '<a href="https://colab.research.google.com/#fileId=https://github.com/AIRI-Institute/StyleFeatureEditor/blob/main/notebook/StyleFeatureEditor_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>' |
| | '</div>' |
| | ) |
| | with gr.Row(): |
| | with gr.Column(): |
| | with gr.Accordion("Input Image", open=True): |
| | input_image = gr.Image(label="Input image you want to edit", type="pil", height=300) |
| | align = gr.Checkbox(label="Align (crop and resize) the input image. For SFE to work well, it is necessary to align the input if it is not.", value=True) |
| | with gr.Accordion("Predefined Editings", open=True): |
| | with gr.Accordion("Description", open=False): |
| | gr.Markdown('''A branch of predefined editings gained from InterfaceGAN, Stylespace, GANSpace and StyleClip mappers. Look at the table below to see which direction is responsible for which editings. |
| | |
| | **Editing power** -- the greater the absolute value of this parameter, the more the selected edit will appear. Better use values in the range 7 - 13, lower values may not give the desired edit, higher values -- on the contrary -- may apply edit too much and create artefacts. |
| | |
| | **Positive effect** -- the effect applied to the image when positive editing power is used. |
| | |
| | **Negative effect** -- the effect applied to the image when negative editing power is used. It is usually the opposite of the positive effect. |
| | ''' |
| | ) |
| |
|
| | gr.Dataframe(value=editings_table, datatype=["markdown","markdown","markdown","markdown"], interactive=False, wrap=True, |
| | column_widths=["25px", "25px", "25px", "25px"], height=300) |
| | with gr.Row(): |
| | predef_editing_direction = gr.Dropdown(list(PREDEFINED_EDITINGS_DATA.keys()), label="Editing direction", value="smile") |
| | predef_editing_power = gr.Slider(-20, 20, value=7, step=0.1, label="Editing power") |
| | btn_predef = gr.Button("Edit image") |
| |
|
| | with gr.Accordion("Text Prompt Editings", open=False): |
| | with gr.Accordion("Description", open=False): |
| | gr.Markdown('''You can alse use editings from text prompts via **StyleClip Global Mapper** (https://arxiv.org/abs/2103.17249) or **DeltaEdit** (https://arxiv.org/abs/2303.06285). You just need to choose: |
| | **Method** -- method to use, StyleClip or DeltaEdit |
| | |
| | **Editing power** -- the greater the absolute value of this parameter, the more the selected edit will appear. |
| | |
| | **Neutral prompt** -- some neutral description of the original image (e.g. "a face"). |
| | |
| | **Target prompt** -- text that contains the desired edit (e.g. "a smilling face"). |
| | |
| | **Disentanglement** -- positive number, the less this attribute -- the more related attributes will also be changed (e.g. for grey hair editing, wrinkle, skin colour and glasses may also be edited) |
| | ''') |
| | edit_method = gr.Dropdown(["StyleClip", "DeltaEdit"], label="Editing method", value="StyleClip") |
| | neutral_prompt = gr.Textbox(value="face with hair", label="Neutreal prompt (e.g. 'a face')") |
| | target_prompt = gr.Textbox(value="face with fire hair", label="Target prompt (e.g. 'a smilling face')") |
| | styleclip_editing_power = gr.Slider(-50, 50, value=10, step=1, label="Editing power") |
| | disentanglement = gr.Slider(0, 1, value=0.1, step=0.01, label="Disentanglement") |
| | btn_clip = gr.Button("Edit image") |
| |
|
| | with gr.Accordion("Mask settings (optional)", open=False): |
| | gr.Markdown('''If some artefacts appear during editing (or some details disappear), you can specify an image mask to select which regions of the image should not be edited. The mask must have a size of 1024 x 1024 and represent an inversion of the original image. |
| | |
| | ''' |
| | ) |
| | mask = gr.Image(label="Upload mask for editing", type="pil", height=350) |
| | with gr.Accordion("Mask generating", open=False): |
| | gr.Markdown("Here you can generate mask that separates face (with hair) from the background.") |
| | with gr.Row(): |
| | input_mask = gr.Image(label="Input image for mask generating", type="pil", height=240) |
| | output_mask = gr.Image(label="Generated mask", height=240) |
| | error_message_mask = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message") |
| | align_mask = gr.Checkbox(label="To align (crop and resize image) or not. Only uncheck this box if the original image has already been aligned.", value=True) |
| | mask_trashhold = gr.Slider(0, 1, value=0.9, step=0.001, label="Mask trashold", |
| | info="The more this parameter, the more is face part, and the less is background part.") |
| | btn_mask = gr.Button("Generate mask") |
| |
|
| | with gr.Column(): |
| | with gr.Row(): |
| | output_align = gr.Image(label="Alignet original image", visible=True) |
| | output_unalign = gr.Image(label="Unalinget editing result", visible=True) |
| | with gr.Row(): |
| | output_inv = gr.Image(label="Inversion result", visible=True) |
| | output_edit = gr.Image(label="Editing result", visible=True) |
| | error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message") |
| | gr.Markdown("If artefacts appear during editing -- try lowering the editing power or using a mask.") |
| | gr.Examples( |
| | label="Input Examples for editing", |
| | examples=[ |
| | ["images/scarlet.jpg"], |
| | ["images/gosling.jpg"], |
| | ["images/robert.png"], |
| | ["images/smith.jpg"], |
| | ["images/watson.jpeg"], |
| | ], |
| | inputs=[input_image], |
| | examples_per_page=5 |
| | ) |
| | gr.Examples( |
| | label="Mask Examples for editing", |
| | examples=[ |
| | ["images/scarlet_mask.webp"], |
| | ["images/gosling_mask.webp"], |
| | ["images/robert_mask.webp"], |
| | ["images/smith_mask.webp"], |
| | ["images/watson_mask.webp"], |
| | ], |
| | inputs=[mask] |
| | ) |
| | gr.Examples( |
| | label="Input Examples for Mask generation", |
| | examples=[ |
| | ["images/scarlet.jpg"], |
| | ["images/gosling.jpg"], |
| | ["images/robert.png"], |
| | ["images/smith.jpg"], |
| | ["images/watson.jpeg"], |
| | ], |
| | inputs=[input_mask] |
| | ) |
| |
|
| |
|
| | btn_predef.click( |
| | fn=edit_image, |
| | inputs=[input_image, predef_editing_direction, predef_editing_power, align, mask], |
| | outputs=[output_edit, output_inv, error_message, output_align, output_unalign] |
| | ) |
| | btn_clip.click( |
| | fn=edit_image_clip, |
| | inputs=[input_image, neutral_prompt, target_prompt, disentanglement, styleclip_editing_power, align, mask, edit_method], |
| | outputs=[output_edit, output_inv, error_message, output_align, output_unalign,] |
| | ) |
| | btn_mask.click( |
| | fn=get_mask, |
| | inputs=[input_mask, align_mask, mask_trashhold], |
| | outputs=[output_mask, error_message_mask] |
| | ) |
| |
|
| | gr.Markdown('''To cite the paper by the authors |
| | ``` |
| | @InProceedings{Bobkov_2024_CVPR, |
| | author = {Bobkov, Denis and Titov, Vadim and Alanov, Aibek and Vetrov, Dmitry}, |
| | title = {The Devil is in the Details: StyleFeatureEditor for Detail-Rich StyleGAN Inversion and High Quality Image Editing}, |
| | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, |
| | month = {June}, |
| | year = {2024}, |
| | pages = {9337-9346} |
| | } |
| | ``` |
| | ''') |
| | return demo |
| |
|
| |
|
| | if __name__ == "__main__": |
| | demo = get_demo() |
| | demo.launch(server_name="0.0.0.0", server_port=7860) |
| |
|
