app.py

import cv2
import torch
import random
import numpy as np

import spaces

import PIL
from PIL import Image
from typing import Tuple, List

import diffusers
from diffusers.utils import load_image
from diffusers.models import ControlNetModel
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel

from huggingface_hub import hf_hub_download

from insightface.app import FaceAnalysis

from style_template import styles
from pipeline_stable_diffusion_xl_instantid_full import StableDiffusionXLInstantIDPipeline, draw_kps

import gradio as gr

from depth_anything.dpt import DepthAnything
from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet

import torch.nn.functional as F
from torchvision.transforms import Compose

# global variable
MAX_SEED = np.iinfo(np.int32).max
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "(No style)"
enable_lcm_arg = False

# download checkpoints
from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="Super-shuhe/InstantID-FaceID-6M", filename="controlnet/config.json", local_dir="./checkpoints")
hf_hub_download(
    repo_id="Super-shuhe/InstantID-FaceID-6M",
    filename="controlnet/diffusion_pytorch_model.safetensors",
    local_dir="./checkpoints",
)
hf_hub_download(repo_id="Super-shuhe/InstantID-FaceID-6M", filename="pytorch_model.bin", local_dir="./checkpoints")

# Load face encoder
app = FaceAnalysis(
    name="antelopev2",
    root="./",
    providers=["CPUExecutionProvider"],
)
app.prepare(ctx_id=0, det_size=(640, 640))

depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(device).eval()

transform = Compose([
    Resize(
        width=518,
        height=518,
        resize_target=False,
        keep_aspect_ratio=True,
        ensure_multiple_of=14,
        resize_method='lower_bound',
        image_interpolation_method=cv2.INTER_CUBIC,
    ),
    NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    PrepareForNet(),
])

# Path to InstantID models
face_adapter = f"./checkpoints/pytorch_model.bin"
controlnet_path = f"./checkpoints/controlnet"

# Load pipeline face ControlNetModel
controlnet_identitynet = ControlNetModel.from_pretrained(
    controlnet_path, torch_dtype=dtype
)

# controlnet-canny/depth
controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"

controlnet_canny = ControlNetModel.from_pretrained(
    controlnet_canny_model, torch_dtype=dtype
).to(device)
controlnet_depth = ControlNetModel.from_pretrained(
    controlnet_depth_model, torch_dtype=dtype
).to(device)

def get_depth_map(image):
    
    image = np.array(image) / 255.0

    h, w = image.shape[:2]

    image = transform({'image': image})['image']
    image = torch.from_numpy(image).unsqueeze(0).to("cuda")

    with torch.no_grad():
        depth = depth_anything(image)

    depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
    depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0

    depth = depth.cpu().numpy().astype(np.uint8)

    depth_image = Image.fromarray(depth)

    return depth_image

def get_canny_image(image, t1=100, t2=200):
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    edges = cv2.Canny(image, t1, t2)
    return Image.fromarray(edges, "L")

controlnet_map = {
    "canny": controlnet_canny,
    "depth": controlnet_depth,
}
controlnet_map_fn = {
    "canny": get_canny_image,
    "depth": get_depth_map,
}

pretrained_model_name_or_path = "SG161222/RealVisXL_V5.0"

pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
    pretrained_model_name_or_path,
    controlnet=[controlnet_identitynet],
    torch_dtype=dtype,
    safety_checker=None,
    feature_extractor=None,
).to(device)

pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
    pipe.scheduler.config
)

# load and disable LCM
pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
pipe.disable_lora()

pipe.cuda()
pipe.load_ip_adapter_instantid(face_adapter)
pipe.image_proj_model.to("cuda")
pipe.unet.to("cuda")

def toggle_lcm_ui(value):
    if value:
        return (
            gr.update(minimum=0, maximum=100, step=1, value=5),
            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=1.5),
        )
    else:
        return (
            gr.update(minimum=5, maximum=100, step=1, value=30),
            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=5),
        )

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

def remove_tips():
    return gr.update(visible=False)

def convert_from_cv2_to_image(img: np.ndarray) -> Image:
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

def convert_from_image_to_cv2(img: Image) -> np.ndarray:
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def resize_img(
    input_image,
    max_side=1280,
    min_side=1024,
    size=None,
    pad_to_max_side=False,
    mode=PIL.Image.BILINEAR,
    base_pixel_number=64,
):
    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        ratio = min_side / min(h, w)
        w, h = round(ratio * w), round(ratio * h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[
            offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new
        ] = np.array(input_image)
        input_image = Image.fromarray(res)
    return input_image

def apply_style(
    style_name: str, positive: str, negative: str = ""
) -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + " " + negative

def update_face_gallery(files):
    return gr.update(value=files, visible=True)

@spaces.GPU
def generate_image(
    face_images_path,  # Now accepts a list of image paths
    pose_image_path,
    prompt,
    negative_prompt,
    style_name,
    num_steps,
    identitynet_strength_ratio,
    adapter_strength_ratio,
    canny_strength,
    depth_strength,
    controlnet_selection,
    guidance_scale,
    seed,
    scheduler,
    enable_LCM,
    enhance_face_region,
    progress=gr.Progress(track_tqdm=True),
):

    if enable_LCM:
        pipe.scheduler = diffusers.LCMScheduler.from_config(pipe.scheduler.config)
        pipe.enable_lora()
    else:
        pipe.disable_lora()
        scheduler_class_name = scheduler.split("-")[0]

        add_kwargs = {}
        if len(scheduler.split("-")) > 1:
            add_kwargs["use_karras_sigmas"] = True
        if len(scheduler.split("-")) > 2:
            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
        scheduler = getattr(diffusers, scheduler_class_name)
        pipe.scheduler = scheduler.from_config(pipe.scheduler.config, **add_kwargs)

    if face_images_path is None or len(face_images_path) == 0:
        raise gr.Error(
            f"Cannot find any input face images! Please upload at least one face image"
        )

    if prompt is None:
        prompt = "a person"

    # apply the style template
    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)

    # Use the first face image for face keypoints and size reference
    reference_face_path = face_images_path[0] if isinstance(face_images_path, list) else face_images_path
    reference_face_image = load_image(reference_face_path)
    reference_face_image = resize_img(reference_face_image, max_side=1024)
    reference_face_cv2 = convert_from_image_to_cv2(reference_face_image)
    height, width, _ = reference_face_cv2.shape

    # Initialize a list to collect face embeddings
    face_embeddings = []
    
    # Process each face image if multiple images are provided
    face_image_paths = face_images_path if isinstance(face_images_path, list) else [face_images_path]
    
    for face_path in face_image_paths:
        face_img = load_image(face_path)
        face_img = resize_img(face_img, max_side=1024)
        face_img_cv2 = convert_from_image_to_cv2(face_img)
        
        # Extract face features
        face_info = app.get(face_img_cv2)
        
        if len(face_info) == 0:
            print(f"Warning: Unable to detect a face in {face_path}. Skipping this image.")
            continue
            
        # Use the largest face in each image
        face_info = sorted(
            face_info,
            key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
        )[-1]
        
        # Collect the embedding
        face_embeddings.append(torch.tensor(face_info["embedding"]).unsqueeze(0))
    
    if len(face_embeddings) == 0:
        raise gr.Error(
            f"Unable to detect a face in any of the uploaded images. Please upload different photos with clear faces."
        )
    
    # Average the face embeddings
    if len(face_embeddings) == 1:
        face_emb = face_embeddings[0].squeeze().numpy()  # Use as is if only one image
    else:
        # Stack and compute mean along the batch dimension
        face_emb = torch.mean(torch.cat(face_embeddings, dim=0), dim=0).numpy()
        print(f"Averaged {len(face_embeddings)} face embeddings")

    # Extract keypoints from the reference face for ControlNet
    reference_face_info = app.get(reference_face_cv2)
    if len(reference_face_info) == 0:
        raise gr.Error(
            f"Unable to detect a face in the reference image for keypoints. Please upload a different photo with a clear face."
        )
    reference_face_info = sorted(
        reference_face_info,
        key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
    )[-1]  # Use the largest face
    
    face_kps = draw_kps(convert_from_cv2_to_image(reference_face_cv2), reference_face_info["kps"])
    img_controlnet = reference_face_image
    
    if pose_image_path is not None:
        pose_image = load_image(pose_image_path)
        pose_image = resize_img(pose_image, max_side=1024)
        img_controlnet = pose_image
        pose_image_cv2 = convert_from_image_to_cv2(pose_image)

        face_info = app.get(pose_image_cv2)

        if len(face_info) == 0:
            raise gr.Error(
                f"Cannot find any face in the reference image! Please upload another person image"
            )

        face_info = face_info[-1]
        face_kps = draw_kps(pose_image, face_info["kps"])

        width, height = face_kps.size

    if enhance_face_region:
        control_mask = np.zeros([height, width, 3])
        x1, y1, x2, y2 = face_info["bbox"]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        control_mask[y1:y2, x1:x2] = 255
        control_mask = Image.fromarray(control_mask.astype(np.uint8))
    else:
        control_mask = None

    if len(controlnet_selection) > 0:
        controlnet_scales = {
            "canny": canny_strength,
            "depth": depth_strength,
        }
        pipe.controlnet = MultiControlNetModel(
            [controlnet_identitynet]
            + [controlnet_map[s] for s in controlnet_selection]
        )
        control_scales = [float(identitynet_strength_ratio)] + [
            controlnet_scales[s] for s in controlnet_selection
        ]
        control_images = [face_kps] + [
            controlnet_map_fn[s](img_controlnet).resize((width, height))
            for s in controlnet_selection
        ]
    else:
        pipe.controlnet = controlnet_identitynet
        control_scales = float(identitynet_strength_ratio)
        control_images = face_kps

    generator = torch.Generator(device=device).manual_seed(seed)

    print("Start inference...")
    print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")

    pipe.set_ip_adapter_scale(adapter_strength_ratio)
    images = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image_embeds=face_emb,
        image=control_images,
        control_mask=control_mask,
        controlnet_conditioning_scale=control_scales,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        height=height,
        width=width,
        generator=generator,
    ).images

    return images[0], gr.update(visible=True)

def get_example():
    case = [
        [
            "./examples/yann-lecun_resize.jpg",
            None,
            "a man",
            "Spring Festival",
            "(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, photo, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green",
        ],
        # Add more examples as needed
    ]
    return case

def run_for_examples(face_file, pose_file, prompt, style, negative_prompt):
    return generate_image(
        face_file,
        pose_file,
        prompt,
        negative_prompt,
        style,
        20,  # num_steps
        0.8,  # identitynet_strength_ratio
        0.8,  # adapter_strength_ratio
        0.3,  # canny_strength
        0.5,  # depth_strength
        ["depth", "canny"],  # controlnet_selection
        5.0,  # guidance_scale
        42,  # seed
        "EulerDiscreteScheduler",  # scheduler
        False,  # enable_LCM
        True,  # enable_Face_Region
    )

# Description
title = r"""
<h1 align="center">InstantID: Zero-shot Identity-Preserving Generation with Multi-Face Averaging</h1>
"""

article = r"""
---
📝 **Citation**
<br>
If our work is helpful for your research or applications, please cite us via:
```bibtex
@article{wang2024instantid,
  title={InstantID: Zero-shot Identity-Preserving Generation in Seconds},
  author={Wang, Qixun and Bai, Xu and Wang, Haofan and Qin, Zekui and Chen, Anthony},
  journal={arXiv preprint arXiv:2401.07519},
  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>.
"""

tips = r"""
### Usage tips of InstantID with Multi-Face Averaging
1. Upload multiple photos of the same person for better identity preservation through face embedding averaging.
2. If you're not satisfied with the similarity, try increasing the weight of "IdentityNet Strength" and "Adapter Strength."    
3. If you feel that the saturation is too high, first decrease the Adapter strength. If it remains too high, then decrease the IdentityNet strength.
4. If you find that text control is not as expected, decrease Adapter strength.
5. If you find that realistic style is not good enough, go for our Github repo and use a more realistic base model.
"""

css = """
.gradio-container {width: 85% !important}
"""
with gr.Blocks(css=css) as demo:
    # description
    gr.Markdown(title)

    with gr.Row():
        with gr.Column():
            with gr.Row(equal_height=True):
                # Change from single image to multiple files
                face_files = gr.Files(
                    label="Upload photos of your face (1 or more)",
                    file_types=["image"]
                )
            
            face_gallery = gr.Gallery(
                label="Your uploaded face images",
                visible=True,
                columns=5,
                rows=1,
                height=150
            )

            # prompt
            prompt = gr.Textbox(
                label="Prompt",
                info="Give simple prompt is enough to achieve good face fidelity",
                placeholder="A photo of a person",
                value="",
            )

            submit = gr.Button("Submit", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                enable_LCM = gr.Checkbox(
                    label="Enable Fast Inference with LCM", value=enable_lcm_arg,
                    info="LCM speeds up the inference step, the trade-off is the quality of the generated image. It performs better with portrait face images rather than distant faces",
                )
                style = gr.Dropdown(
                    label="Style template",
                    choices=STYLE_NAMES,
                    value=DEFAULT_STYLE_NAME,
                )

                # strength
                identitynet_strength_ratio = gr.Slider(
                    label="IdentityNet strength (for fidelity)",
                    minimum=0,
                    maximum=1.5,
                    step=0.05,
                    value=0.80,
                )
                adapter_strength_ratio = gr.Slider(
                    label="Image adapter strength (for detail)",
                    minimum=0,
                    maximum=1.5,
                    step=0.05,
                    value=0.80,
                )
            with gr.Accordion("Controlnet", open=False):
                # optional: upload a reference pose image
                pose_file = gr.Image(
                    label="Upload a reference pose image (Optional)",
                    type="filepath",
                )
                controlnet_selection = gr.CheckboxGroup(
                    ["canny", "depth"], label="Controlnet", value=[],
                    info="Use canny for edge detection, and depth for depth map estimation to control the generation process"
                )
                canny_strength = gr.Slider(
                    label="Canny strength",
                    minimum=0,
                    maximum=1.5,
                    step=0.05,
                    value=0.3,
                )
                depth_strength = gr.Slider(
                    label="Depth strength",
                    minimum=0,
                    maximum=1.5,
                    step=0.05,
                    value=0.5,
                )
            with gr.Accordion(open=False, label="Advanced Options"):
                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    placeholder="low quality",
                    value="(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green",
                )
                num_steps = gr.Slider(
                    label="Number of sample steps",
                    minimum=1,
                    maximum=100,
                    step=1,
                    value=5 if enable_lcm_arg else 30,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=20.0,
                    step=0.1,
                    value=0.0 if enable_lcm_arg else 5.0,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42,
                )
                schedulers = [
                    "DEISMultistepScheduler",
                    "HeunDiscreteScheduler",
                    "EulerDiscreteScheduler",
                    "DPMSolverMultistepScheduler",
                    "DPMSolverMultistepScheduler-Karras",
                    "DPMSolverMultistepScheduler-Karras-SDE",
                ]
                scheduler = gr.Dropdown(
                    label="Schedulers",
                    choices=schedulers,
                    value="EulerDiscreteScheduler",
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                enhance_face_region = gr.Checkbox(label="Enhance non-face region", value=True)

        with gr.Column(scale=1):
            gallery = gr.Image(label="Generated Images")
            usage_tips = gr.Markdown(
                label="InstantID Usage Tips", value=tips, visible=False
            )

        # Connect file uploads to update the gallery
        face_files.upload(fn=update_face_gallery, inputs=face_files, outputs=face_gallery)

        submit.click(
            fn=remove_tips,
            outputs=usage_tips,
        ).then(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
        ).then(
            fn=generate_image,
            inputs=[
                face_files,  # Changed from face_file to face_files
                pose_file,
                prompt,
                negative_prompt,
                style,
                num_steps,
                identitynet_strength_ratio,
                adapter_strength_ratio,
                canny_strength,
                depth_strength,
                controlnet_selection,
                guidance_scale,
                seed,
                scheduler,
                enable_LCM,
                enhance_face_region,
            ],
            outputs=[gallery, usage_tips],
        )

        enable_LCM.input(
            fn=toggle_lcm_ui,
            inputs=[enable_LCM],
            outputs=[num_steps, guidance_scale],
            queue=False,
        )

    gr.Examples(
        examples=get_example(),
        inputs=[face_files, pose_file, prompt, style, negative_prompt],
        fn=run_for_examples,
        outputs=[gallery, usage_tips],
        cache_examples=True,
    )

    gr.Markdown(article)

demo.queue(api_open=False)
demo.launch()