app.py

import sys
sys.path.append('./')
import gradio as gr
import random
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
from models import iresnet
from sixdrepnet.model import SixDRepNet
import pixel_generator.vec2face.model_vec2face as model_vec2face
MAX_SEED = np.iinfo(np.int32).max
import torch
from time import time


device = "cuda"

def clear_image():
    return None


def clear_generation_time():
    return ""


def generating():
    return "Generating images..."


def done():
    return "Done!"


def sample_nearby_vectors(base_vector, epsilons=[0.3, 0.5, 0.7], percentages=[0.4, 0.4, 0.2]):
    row, col = base_vector.shape
    norm = torch.norm(base_vector, 2, 1, True)
    diff = []
    for i, eps in enumerate(epsilons):
        diff.append(np.random.normal(0, eps, (int(row * percentages[i]), col)))
    diff = np.vstack(diff)
    np.random.shuffle(diff)
    diff = torch.tensor(diff)
    generated_samples = base_vector + diff
    generated_samples = generated_samples / torch.norm(generated_samples, 2, 1, True) * norm
    return generated_samples


def initialize_models():
    pose_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/6DRepNet_300W_LP_AFLW2000.pth", local_dir="./")
    id_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/arcface-r100-glint360k.pth", local_dir="./")
    quality_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/magface-r100-glint360k.pth", local_dir="./")
    generator_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/vec2face_generator.pth", local_dir="./")
    generator = model_vec2face.__dict__["vec2face_vit_base_patch16"](mask_ratio_mu=0.15, mask_ratio_std=0.25,
                                                mask_ratio_min=0.1, mask_ratio_max=0.5,
                                                use_rep=True,
                                                rep_dim=512,
                                                rep_drop_prob=0.,
                                                use_class_label=False)
    generator = generator.to(device)
    checkpoint = torch.load(generator_weights, map_location=device)
    generator.load_state_dict(checkpoint['model_vec2face'])
    generator.eval()

    id_model = iresnet("100", fp16=True).to(device)
    id_model.load_state_dict(torch.load(id_model_weights, map_location=device))
    id_model.eval()

    quality_model = iresnet("100", fp16=True).to(device)
    quality_model.load_state_dict(torch.load(quality_model_weights, map_location=device))
    quality_model.eval()

    pose_model = SixDRepNet(backbone_name='RepVGG-B1g2',
                            backbone_file='',
                            deploy=True,
                            pretrained=False
                            ).to(device)
    pose_model.load_state_dict(torch.load(pose_model_weights))
    pose_model.eval()

    return generator, id_model, pose_model, quality_model


def image_generation(input_image, quality, use_target_pose, pose, dimension, progress=gr.Progress()):
    generator, id_model, pose_model, quality_model = initialize_models()

    generated_images = []
    if input_image is None:
        feature = np.random.normal(0, 1.0, (1, 512))
    else:
        input_image = np.transpose(input_image, (2, 0, 1))
        input_image = torch.from_numpy(input_image).unsqueeze(0).float().to(device)
        input_image.div_(255).sub_(0.5).div_(0.5)
        feature = id_model(input_image).clone().detach().cpu().numpy()

    if not use_target_pose:
        features = []
        norm = np.linalg.norm(feature, 2, 1, True)
        for i in progress.tqdm(np.arange(0, 4.8, 2), desc="Generating images"):
            updated_feature = feature
            updated_feature[0][dimension] = feature[0][dimension] + i
            updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
            features.append(updated_feature)
        features = torch.tensor(np.vstack(features)).float().to(device)
        if quality > 25:
            images, _ = generator.gen_image(features, quality_model, id_model, q_target=quality)
        else:
            _, _, images, *_ = generator(features)
    else:
        features = torch.repeat_interleave(torch.tensor(feature), 3, dim=0)
        features = sample_nearby_vectors(features, [0.7], [1]).float().to(device)
        if quality > 25 or pose > 20:
            images, _ = generator.gen_image(features, quality_model, id_model, pose_model=pose_model,
                                            q_target=quality, pose=pose, class_rep=features)
        else:
            _, _, images, *_ = generator(features)

    images = ((images.permute(0, 2, 3, 1).clip(-1, 1).detach().cpu().numpy() + 1) / 2 * 255).astype(np.uint8)
    for image in progress.tqdm(images, desc="Processing images"):
        generated_images.append(Image.fromarray(image))

    return generated_images


def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose, progress=gr.Progress()):
    # Ensure all dimension numbers are within [0, 512)
    num1, num2, num3, num4 = [max(0, min(int(n), 511)) for n in [num1, num2, num3, num4]]

    # Use the provided random seed
    random.seed(random_seed)
    np.random.seed(random_seed)
    if image_input is None:
        input_data = None
    else:
        # Process the uploaded image
        input_data = Image.open(image_input)
        input_data = np.array(input_data.resize((112, 112)))

    generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4], progress)

    return generated_images


def select_image(value, images):
    # Convert the float value (0 to 4) to an integer index (0 to 9)
    index = int(value / 2)
    return images[index]

def toggle_inputs(use_pose):
    return [
        gr.update(visible=use_pose, interactive=use_pose),  # target_pose
        gr.update(interactive=not use_pose),  # num1
        gr.update(interactive=not use_pose),  # num2
        gr.update(interactive=not use_pose),  # num3
        gr.update(interactive=not use_pose),  # num4
    ]


# 4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
def main():
    with gr.Blocks() as demo:
        title = r"""

            <h1 align="center">Vec2Face: Scaling Face Dataset Generation with Loosely Constrained Vectors</h1>

            """

        description = r"""

            <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HaiyuWu/vec2face' target='_blank'><b>Vec2Face: Scaling Face Dataset Generation with Loosely Constrained Vectors</b></a>.<br>



            How to use:<br>

            1. Upload an image with a cropped face image or directly click <b>Submit</b> button, three images will be shown on the right. 

            2. You can control the image quality, image pose, and modify the values in the target dimensions to change the output images. 

            3. The output results will shown three results of dimension modification or pose images.

            4. Enjoy! 😊

            """

        gr.Markdown(title)
        gr.Markdown(description)
        with gr.Row():
            with gr.Column():
                image_file = gr.Image(label="Upload an image (optional)", type="filepath")

                gr.Markdown("""

                ## Dimension Modification

                Enter the values for the dimensions you want to modify (0-511). 

                """)

                with gr.Row():
                    num1 = gr.Number(label="Dimension 1", value=0, minimum=0, maximum=511, step=1)
                    num2 = gr.Number(label="Dimension 2", value=50, minimum=0, maximum=511, step=1)
                    num3 = gr.Number(label="Dimension 3", value=100, minimum=0, maximum=511, step=1)
                    num4 = gr.Number(label="Dimension 4", value=200, minimum=0, maximum=511, step=1)

                random_seed = gr.Number(label="Random Seed", value=42, minimum=0, maximum=MAX_SEED, step=1)
                target_quality = gr.Slider(label="Minimum Quality", minimum=22, maximum=35, step=1, value=24)

                with gr.Row():
                    use_target_pose = gr.Checkbox(label="Use Target Pose")
                    target_pose = gr.Slider(label="Target Pose", value=0, minimum=0, maximum=90, step=1, visible=False)

                submit = gr.Button("Submit", variant="primary")

                gr.Markdown("""

                            ## Usage tips of Vec2Face

                            - Directly clicking "Submit" button will give you results from a randomly sampled vector. 

                            - If you want to modify more dimensions, please write your own code. Code snippets in [Vec2Face repo](https://github.com/HaiyuWu/vec2face) might be helpful.

                            - If you want to create extreme pose image (e.g., >70), please do not set image quality larger than 27.

                            - <span style="color: red;">!</span> <span style="color: red;">!</span> <span style="color: red;">!</span> **Due to the limitation of SixDRepNet (pose estimator), pose editing results might be corrupted/incorrect. For better performance, you can integrade other pose estimators.** <span style="color: red;">!</span> <span style="color: red;">!</span> <span style="color: red;">!</span>

                            - For better experience, we suggest you to run code on a GPU machine.

                            """)

            with gr.Column():
                gallery = gr.Image(label="Generated Image")
                generation_time = gr.Textbox(label="Generation Status")
                incremental_value_slider = gr.Slider(
                    label="Result of dimension modification or results of pose images",
                    minimum=0, maximum=4, step=2, value=0
                )
                gr.Markdown("""

                            - These values are added to the dimensions (before normalization), **please ignore it if pose editing is on**.

                            """)

        use_target_pose.change(
            fn=toggle_inputs,
            inputs=[use_target_pose],
            outputs=[target_pose, num1, num2, num3, num4]
        )

        generated_images = gr.State([])

        submit.click(
            fn=clear_image,
            inputs=[],
            outputs=[gallery]
        ).then(
            fn=generating,
            inputs=[],
            outputs=[generation_time]
        ).then(
            fn=process_input,
            inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
            outputs=[generated_images]
        ).then(
            fn=done,
            inputs=[],
            outputs=[generation_time]
        ).then(
            fn=select_image,
            inputs=[incremental_value_slider, generated_images],
            outputs=[gallery]
        )
        # submit.click(
        #     fn=process_input,
        #     inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
        #     outputs=[generated_images]
        # ).then(
        #     fn=select_image,
        #     inputs=[incremental_value_slider, generated_images],
        #     outputs=[gallery]
        # )

        incremental_value_slider.change(
            fn=select_image,
            inputs=[incremental_value_slider, generated_images],
            outputs=[gallery]
        )
        article = r"""

        ---

        📝 **Citation**

        <br>

        If our work is helpful for your research or applications, please cite us via:

        ```bibtex

        @article{wu2024vec2face,

        title={Vec2Face: Scaling Face Dataset Generation with Loosely Constrained Vectors},

        author={Wu, Haiyu and Singh, Jaskirat and Tian, Sicong and Zheng, Liang and Bowyer, Kevin W.},

        year={2024}

        }

        ```

        📧 **Contact**

        <br>

        If you have any questions, please feel free to open an issue or directly reach us out at <b>[email protected]</b>.

        """
        gr.Markdown(article)

    demo.launch(share=True)


if __name__ == "__main__":
    main()