app.py

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)): # output_align, output_unalign
    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) # 100
                    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)