app.py

# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import gradio as gr
import huggingface_hub
import pillow_avif
import spaces
import gc
from huggingface_hub import snapshot_download
from pillow_heif import register_heif_opener
from PIL import Image, ImageDraw, ImageFont

from pathlib import Path
import numpy as np
import cv2
import tensorflow as tf
from mtcnn import MTCNN
from insightface.utils import face_align
import facexlib
import torch
from modules.inferencer import IDPatchInferencer
from rtmlib import Body
from utils.draw_condition import draw_openpose_from_mmpose

# Register HEIF support for Pillow
register_heif_opener()

loaded_pipeline_config = {
    'pipeline': None,
    'face_encoder': None,
    'face_detector': None
}

body_estimator = Body(to_openpose=False, mode='balanced', backend='onnxruntime', device='cpu')

def pil_to_cv2(pil_image):
    """PIL.Image -> OpenCV BGR Image"""
    cv2_image = np.array(pil_image)
    cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_RGB2BGR)
    return cv2_image

def mtcnn_to_kps(mtcnn_results):
    kps = np.array([mtcnn_results[0]['keypoints']['left_eye'], mtcnn_results[0]['keypoints']['right_eye'], mtcnn_results[0]['keypoints']['nose'], mtcnn_results[0]['keypoints']['mouth_left'], mtcnn_results[0]['keypoints']['mouth_right']])
    return kps

def extract_face_emb(arcface_encoder, cropped_face):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    face_image = torch.from_numpy(cropped_face).unsqueeze(0).permute(0,3,1,2) / 255.
    face_image = 2 * face_image - 1
    face_image = face_image.to(device).contiguous()
    face_emb = arcface_encoder(face_image)[0]
    return face_emb


def download_models():
    snapshot_download(repo_id='ByteDance/ID-Patch', revision="5e5434dc43a8d1325aade8b0da65d96d7c4cf3d9", local_dir='./models/ID-Patch', local_dir_use_symlinks=False)
    snapshot_download(repo_id='RunDiffusion/Juggernaut-X-v10', revision="main", local_dir='./models/Juggernaut-X-v10', local_dir_use_symlinks=False)


def init_pipeline():
    pipeline = loaded_pipeline_config['pipeline']
    gc.collect()

    model_path = f'./models/ID-Patch'
    print(f'loading model from {model_path}')

    pipeline = IDPatchInferencer(base_model_path='./models/Juggernaut-X-v10', idp_model_path='./models/ID-Patch')
    
    loaded_pipeline_config['pipeline'] = pipeline
    return pipeline

# Future works: Add more model variants 
def prepare_pipeline():
    pipeline = loaded_pipeline_config['pipeline']
    return pipeline


def add_safety_watermark(image, text='AI Generated: ID-Patch', font_path=None):
    width, height = image.size
    draw = ImageDraw.Draw(image)

    font_size = int(height * 0.028)
    if font_path:
        font = ImageFont.truetype(font_path, font_size)
    else:
        font = ImageFont.load_default(size=font_size)
    
    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
    x = width - text_width - 10
    y = height - text_height - 20

    shadow_offset = 2
    shadow_color = "black"
    draw.text((x + shadow_offset, y + shadow_offset), text, font=font, fill=shadow_color)

    draw.text((x, y), text, font=font, fill="white")

    return image


@spaces.GPU(duration=60)
def generate_image(
    id_images, 
    id_order,
    control_image, 
    prompt, 
    seed, 
    guidance_scale, 
    num_steps, 
    controlnet_conditioning_scale,
    id_injection_ratio,
    negative_prompt
):
    try:
        print("======= Start Generating =======")
        print(f"ID Images uploaded: {len(id_images)}")

        id_images_pil = []
        for idx, img in enumerate(id_images):
            img = img[0]
            print(f"ID Image {idx} size: {img.size}")
            id_images_pil.append(img)
        id_images = id_images_pil

        if id_order is not None and id_order.strip() != "":
            id_order = id_order.split(',')
            sorted_id_images = [id_images[int(i)] for i in id_order]
            id_images = sorted_id_images
        else:
            id_order = [i for i in range(len(id_images))]

        print(f"Control Image size: {control_image.size}")

        pipeline = prepare_pipeline()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        arcface_encoder = facexlib.recognition.init_recognition_model('arcface', device=device)
        tf.config.set_visible_devices([], 'GPU')
        mtcnn_inferencer = MTCNN() # MTCNN might be slow, could be replaced by other face detectors, as long as it provides 5 keypoints

        if seed == 0:
            seed = torch.seed() & 0xFFFFFFFF

        face_embs = []
        for subject in id_images:
            image_subject = pil_to_cv2(subject)

            mtcnn_subject = mtcnn_inferencer.detect_faces(image_subject[:,:,::-1])
            if not mtcnn_subject:
                print("Warning: No face detected in uploaded identity image.")
                continue  # skip this image

            try:
                kps_subject = mtcnn_to_kps(mtcnn_subject)
                cropped_face_subject = face_align.norm_crop(image_subject, landmark=kps_subject, image_size=112)
                emb = extract_face_emb(arcface_encoder, cropped_face_subject)
                face_embs.append(emb)
            except Exception as e:
                print(f"Error processing face: {e}")
                continue

        if len(face_embs) == 0:
            raise ValueError("No valid face embeddings extracted. Please upload clear identity images.")
        
        face_embs = torch.stack(face_embs)

        # load pose
        image_reference = pil_to_cv2(control_image)

        # estimate pose
        keypoints, scores = body_estimator(image_reference)
        
        # Check 
        print(f"Keypoints raw output: {keypoints}")
        if keypoints is None:
            raise ValueError("Keypoints is None.")
        if not isinstance(keypoints, (list, np.ndarray)):
            raise ValueError(f"Keypoints type wrong: {type(keypoints)}")
        if len(keypoints) == 0:
            raise ValueError("Keypoints length == 0.")

        keypoints = np.array(keypoints)
        print(f"Keypoints converted to np.array, shape = {keypoints.shape}")

        if len(keypoints.shape) != 3:
            raise ValueError(f"Keypoints wrong shape: {keypoints.shape}")
        if keypoints.shape[0] == 0:
            raise ValueError("No people detected in the pose control image.")

        face_locations = keypoints[:, 0]
        face_locations = sorted(face_locations, key=lambda x: x[0] if isinstance(x, (list, tuple, np.ndarray)) and len(x) > 0 else 0)
        face_locations = torch.from_numpy(np.stack(face_locations))
        
        # Draw OpenPose image
        control_image = Image.fromarray(draw_openpose_from_mmpose(image_reference * 0, keypoints, scores))

        image = pipeline.generate(
            face_embs,
            face_locations,
            control_image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            id_injection_ratio=id_injection_ratio,
            seed=seed
        )
        image = add_safety_watermark(image)
    
    except Exception as e:
        print(e)
        gr.Error(f"An error occurred: {e}")
        return gr.update()

    return gr.update(value=image, label=f"Generated Image, seed = {seed}"), gr.update(value=control_image, label=f"OpenPose")


def generate_examples(id_image_paths, ui_id_order, control_image_path, prompt_text, seed):
    id_images = [Image.open(p).convert('RGB') for p in id_image_paths]
    control_image = Image.open(control_image_path).convert('RGB')
    return generate_image(id_images, ui_id_order, control_image, prompt_text, seed, 5.5, 50, 0.8, 0.8, "nude, worst quality, low quality, normal quality, nsfw, abstract, glitch, deformed, mutated, ugly, disfigured, text, watermark, bad hands, error, jpeg artifacts, blurry, missing fingers")


def load_example(selected_key):
    if selected_key is None:
        return None, None, None, None, None

    example = example_choices[selected_key]
    id_images = [Image.open(p).convert('RGB') for p in example['id_images']]
    control_image = Image.open(example['pose_image']).convert('RGB')
    return (
        id_images,               # For ui_id_image (Gallery)
        example['id_order'],      # For ui_id_order (Textbox)
        control_image,            # For ui_control_image (Image)
        example['prompt'],        # For ui_prompt_text (Textbox)
        example['seed']           # For ui_seed (Number)
    )

# Get all available ID and pose images
man_images = sorted(list(Path('./assets/subjects/man').glob('*.jpg')))
woman_images = sorted(list(Path('./assets/subjects/woman').glob('*.jpg')))
pose_images = sorted(list(Path('./assets/poses').glob('*.png')) + list(Path('./assets/poses').glob('*.jpeg')) + list(Path('./assets/poses').glob('*.jpg')))

def random_select_id_images(num_men, num_women):
    if int(num_men) > len(man_images) or int(num_women) > len(woman_images):
        raise ValueError("Requested more images than available.")
    selected_men = np.random.choice(man_images, size=int(num_men), replace=False)
    selected_women = np.random.choice(woman_images, size=int(num_women), replace=False)
    selected = list(selected_men) + list(selected_women)
    images = [Image.open(p).convert('RGB') for p in selected]
    id_order = ",".join(str(i) for i in range(len(images)))
    return images, id_order




with gr.Blocks() as demo:
    session_state = gr.State({})
    default_model_version = "v1.0"

    gr.HTML("""
    <div style="text-align: center; max-width: 900px; margin: 0 auto;">
        <h1 style="font-size: 1.5rem; font-weight: 700; display: block;">ID-Patch-SDXL</h1>
        <h2 style="font-size: 1.2rem; font-weight: 300; margin-bottom: 1rem; display: block;">Official Gradio Demo for Our CVPR 2025 Paper <br><br>
            <b>"ID-Patch: Robust ID Association for Group Photo Personalization" </b>
        </h2>
        <a href="https://byteaigc.github.io/ID-Patch/">[Project Page]</a>&ensp;
        <a href="https://arxiv.org/abs/2411.13632">[Paper]</a>&ensp;
        <a href="https://damon-demon.github.io/links/ID_Patch_CVPR25_poster.pdf">[Poster]</a>&ensp;
        <a href="https://github.com/bytedance/ID-Patch">[Code]</a>&ensp;
        <a href="https://huggingface.co/ByteDance/ID-Patch">[Model]</a>
    </div>
    """)

    # Add the pipeline image of ID-Patch: assets/pipeline.png and short description
    # with gr.Column(elem_id="pipeline_block"):
    #     gr.Image(
    #         value="./assets/pipeline.png", 
    #         interactive=False, 
    #         show_label=False, 
    #         height=300,
    #         container=False
    #     )
    #     gr.HTML(
    #         """
    #         <div style="text-align:center; font-size:1.2rem; font-weight:300;">
    #             Pipeline of ID-Patch: Build Identity-to-Position Association
    #         </div>
    #         """
    #     )


    gr.Markdown("""
    ### 💡 How to Use This Demo?
    1. **Upload ID images**:  
        - Upload one or more ID images for each person you want to generate.  
            *(The number of uploaded ID images should match the number of people in your pose reference image.)*
    
    2. **ID Order**:
        - List the ID images separated by commas, following the **left-to-right order** of detected faces in the pose reference image.   
            *(ID index starts from 0!)*
        

    3. **Upload a pose reference image**:  
        - Choose an image that shows the desired pose(s) for the people you want to generate.
            *(Tip: If the pose is too complicated, then the face detection and pose detection might fail.)*

    4. **Enter a text prompt**:  
        - Describe the scene you want to create.  
            *(Tip: Try to match the interactions described in your text with the uploaded pose reference.)*

    5. **[Optional] Adjust advanced settings**:  
        Fine-tune generation details if needed.

    6. **Click "Generate"**:  
        Your personalized image will be created. Enjoy!
                
    ### 🔫  Example Playground
    - We offer example settings that users can easily select and load all required settings (identity images, pose image and others) by clicking the **“Load Settings”** button for testing.
                
    - Alternatively, you can randomly sample a specific number of male and female face images from our provided identity image dataset and/or choose a pose from the available options.
    """)
    
    with gr.Row():
        with gr.Column(scale=3):

            example_choices = {
                "Woman Playing Piano (1 People)": {
                    "id_images": ['./assets/subjects/woman/66.jpg'],
                    "id_order": "0",
                    "pose_image": './assets/poses/p1_1.jpeg',
                    "prompt": 'a woman is playing piano, (pianist:1.1), wearing an elegant metallic gold backless gown dress, silver earrings, on the stage, in the spotlight, bright and colorful lighting, LED screen background, vibrant fill light',
                    "seed": 1111
                },
                "Man Playing Piano (1 People)": {
                    "id_images": ['./assets/subjects/man/21.jpg'],
                    "id_order": "0",
                    "pose_image": './assets/poses/p1_2.jpeg',
                    "prompt": 'a man is playing piano, (pianist:1.1), wearing an elegant black havana tuxedo, on the stage, in the spotlight, bright and colorful lighting, LED screen background',
                    "seed": 1111
                },
                "Couple Cheers (2 People)": {
                    "id_images": ['./assets/subjects/man/0.jpg', './assets/subjects/woman/0.jpg'],
                    "id_order": "0,1",
                    "pose_image": './assets/poses/p2.png',
                    "prompt": 'a young couple in front of their burning home still managing to find a moment of joy amidst disaster. cheerfully raise glasses filled with a bright blue drink',
                    "seed": 2222
                },
                "Friend Selfie (3 People)": {
                    "id_images": ['./assets/subjects/woman/26.jpg', './assets/subjects/woman/53.jpg','./assets/subjects/man/52.jpg'],
                    "id_order": "0,1,2",
                    "pose_image": './assets/poses/p3_2.jpeg',
                    "prompt": 'a joyful selfie of three friends, background of television studio setting.',
                    "seed": 3333
                },
                "Happy Piano Moment (4 People)": {
                    "id_images": ['./assets/subjects/man/40.jpg', './assets/subjects/woman/59.jpg','./assets/subjects/woman/79.jpg', './assets/subjects/man/43.jpg'],
                    "id_order": "0,1,2,3",
                    "pose_image": './assets/poses/p4.jpeg',
                    "prompt": 'three adults watch one man playing the piano in a brightly lit, elegant room with vintage decor',
                    "seed": 4444
                },
                "Outdoor Selfie (6 People)": {
                    "id_images": ['./assets/subjects/man/51.jpg', './assets/subjects/man/52.jpg','./assets/subjects/woman/79.jpg', './assets/subjects/woman/66.jpg', './assets/subjects/woman/39.jpg', './assets/subjects/man/49.jpg'],
                    "id_order": "0,1,2,3,4,5",
                    "pose_image": './assets/poses/p6.jpeg',
                    "prompt": 'A joyful group selfie of six adventurous people on a mountain at sunrise. Each person is dressed in outdoor apparel suitable for chilly weather',
                    "seed": 6666
                },
            }

            # Build pose_name_to_path mapping
            pose_name_to_path = {
                example_name: example_data["pose_image"]
                for example_name, example_data in example_choices.items()
            }


            with gr.Accordion("Example Playground", open=False):
                with gr.Column():
                    selected_example = gr.Dropdown(
                        choices=list(example_choices.keys()),
                        label="Example Setting Selections (Identity Images + Pose + Others)",
                        interactive=True
                    )
                    load_example_btn = gr.Button("Load Settings")

                with gr.Row():
                    with gr.Column(scale=3):
                        with gr.Row():
                            num_men_dropdown = gr.Dropdown(
                                choices=[str(i) for i in range(11)], 
                                value="0", 
                                label="Number of Men"
                            )
                            num_women_dropdown = gr.Dropdown(
                                choices=[str(i) for i in range(11)], 
                                value="0", 
                                label="Number of Women"
                            )
                        random_select_button = gr.Button("Random Select Identity Images")

                    with gr.Column(scale=2):
                        pose_dropdown = gr.Dropdown(
                            choices=list(pose_name_to_path.keys()),
                            label="Select Pose Example",
                            interactive=True
                        )
                        pose_select_button = gr.Button("Load Pose")
                
            with gr.Row():
                with gr.Column(scale=3):
                    ui_id_image = gr.Gallery(
                            label="Identity Images", 
                            type="pil", 
                            scale=3, 
                            height=370, 
                            min_width=100, 
                            columns=4, 
                            rows=1, 
                            allow_preview=True, 
                            show_label=True, 
                            interactive=True
                        )


                with gr.Column(scale=2, min_width=100):
                    ui_control_image = gr.Image(label="Pose Reference Image", type="pil", height=370, min_width=100)

            
            ui_prompt_text = gr.Textbox(label="Text Prompt (Describe the image you would like to generate)", value="Portrait, 4K, high quality, cinematic")
            ui_id_order = gr.Textbox(label="ID Order (If not specified, the images will follow the original upload order)", value = None)

                
            ui_btn_generate = gr.Button("Generate")

            with gr.Accordion("Advanced Settings", open=True):
                with gr.Row():
                    ui_num_steps = gr.Number(label="Steps", value=50)
                    ui_seed = gr.Number(label="Seed (0 for random seed)", value=0)
                ui_guidance_scale = gr.Number(label="Guidance Scale", value=5.5, step=0.1)
                ui_controlnet_conditioning_scale = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.05, label="ControlNet Conditioning Scale")
                ui_id_injection_ratio = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.05, label="ID Injection Ratio")
                ui_negative_prompt = gr.Textbox(label="Negative Prompt", value="nude, worst quality, low quality, normal quality, nsfw, abstract, glitch, deformed, mutated, ugly, disfigured, text, watermark, bad hands, error, jpeg artifacts, blurry, missing fingers")



        with gr.Column(scale=2):
            image_output = gr.Image(label="Generated Image", interactive=False, height=615, format='png')
            openpose_control_image = gr.Image(label="OpenPose Image", interactive=False, height=549, format='png')
            


    ui_btn_generate.click(
        generate_image, 
        inputs=[
            ui_id_image, 
            ui_id_order,
            ui_control_image, 
            ui_prompt_text, 
            ui_seed, 
            ui_guidance_scale, 
            ui_num_steps,
            ui_controlnet_conditioning_scale,
            ui_id_injection_ratio,
            ui_negative_prompt 
            
        ], 
        outputs=[image_output, openpose_control_image], 
        concurrency_id="gpu"
    )

    load_example_btn.click(
        load_example,
        inputs=[selected_example],
        outputs=[ui_id_image, ui_id_order, ui_control_image, ui_prompt_text, ui_seed]
    )

    random_select_button.click(
        random_select_id_images,
        inputs=[num_men_dropdown, num_women_dropdown],
        outputs=[ui_id_image, ui_id_order]
    )

    def select_pose_image(pose_name):
        if pose_name not in pose_name_to_path:
            raise ValueError(f"Pose name {pose_name} not found.")
        
        pose_path = pose_name_to_path[pose_name]
        pose_image = Image.open(pose_path).convert('RGB')

        for example_name, example_data in example_choices.items():
            if example_name == pose_name:
                prompt = example_data['prompt']
                break
        else:
            prompt = ""

        return pose_image, prompt
    
    pose_select_button.click(
        select_pose_image,
        inputs=[pose_dropdown],
        outputs=[ui_control_image, ui_prompt_text]
    )

    
    gr.Markdown(
        """
        ---
        ### 📜 Disclaimer and Licenses 
        The images used in this demo are sourced from consented subjects or generated by the models. These pictures are intended solely to show the capabilities of our research. If you have any concerns, please contact us, and we will promptly remove any inappropriate content.
        
        The use of the released code, model, and demo must strictly adhere to the respective licenses. 
        Our code is released under the [Apache License 2.0](https://github.com/bytedance/ID-Patch/blob/main/LICENSE), 
        and our model is released under the [CreativeML Open RAIL++-M License](https://huggingface.co/ByteDance/ID-Patch/blob/main/LICENSE.md) 
        for academic research purposes only. Any manual or automatic downloading of the face models from [InsightFace_Pytorch](https://github.com/TreB1eN/InsightFace_Pytorch), [MTCNN](https://github.com/ipazc/mtcnn), 
        the [Juggernaut-X-v10](https://huggingface.co/RunDiffusion/Juggernaut-X-v10) base model, *etc.*, must follow their original licenses and be used only for academic research purposes.

        This research aims to positively impact the field of Generative AI. Any usage of this method must be responsible and comply with local laws. The developers do not assume any responsibility for any potential misuse. We added the "AI Generated: ID-Patch" watermark for enhanced safety.
        """
    )    

    gr.Markdown(
        """
        ### 📖 Citation

        If you find ID-Patch useful for your research or applications, please cite our paper:

        ```bibtex
        @InProceedings{zhang2025idpatch,
            author    = {Zhang, Yimeng and Zhi, Tiancheng and Liu, Jing and Sang, Shen and Jiang, Liming and Yan, Qing and Liu, Sijia and Luo, Linjie},
            title     = {ID-Patch: Robust ID Association for Group Photo Personalization},
            booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
            month     = {June},
            year      = {2025}
        }
        ```

        We also appreciate it if you could give a star ⭐ to our [Github repository](https://github.com/bytedance/ID-Patch). Thanks a lot!
        """
    )


download_models()
init_pipeline()
demo.launch()