app.py

import cv2, os
import torch
import random
import numpy as np

import spaces

import PIL
from PIL import Image
from typing import Tuple

import diffusers
from diffusers.utils import load_image

from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    UniPCMultistepScheduler,
)

from huggingface_hub import hf_hub_download

from insightface.app import FaceAnalysis

from pipeline_controlnet_xs_sd_xl_instantid import StableDiffusionXLInstantIDXSPipeline, UNetControlNetXSModel

from utils.controlnet_xs import ControlNetXSAdapter
# from controlnet_aux import OpenposeDetector

import gradio as gr

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"
weight_dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32


base_model = 'frankjoshua/realvisxlV40_v40Bakedvae'
vae_path = 'madebyollin/sdxl-vae-fp16-fix'
ckpt = 'RED-AIGC/InstantID-XS'

image_proj_path = os.path.join(ckpt, "image_proj.bin")
cnxs_path =  os.path.join(ckpt, "controlnetxs.bin")
cross_attn_path = os.path.join(ckpt, "cross_attn.bin")


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


def get_ControlNetXS(base_model, cnxs_path, device, size_ratio=0.125, weight_dtype=torch.float16):
    unet = UNet2DConditionModel.from_pretrained(base_model, subfolder="unet").to(device, dtype=weight_dtype)
    controlnet = ControlNetXSAdapter.from_unet(unet, size_ratio=size_ratio, learn_time_embedding=True)
    state_dict = torch.load(cnxs_path, map_location="cpu", weights_only=True)
    ctrl_state_dict = {}
    for key, value in state_dict.items():
        if 'attn2.processor' not in key:
            if 'ctrl_' in key and 'ctrl_to_base' not in key:
                key = key.replace('ctrl_', '')
            if 'up_blocks' in key:
                key = key.replace('up_blocks', 'up_connections')
            ctrl_state_dict[key] = value
    controlnet.load_state_dict(ctrl_state_dict, strict=True)
    controlnet.to(device, dtype=weight_dtype)
    ControlNetXS = UNetControlNetXSModel.from_unet(unet, controlnet).to(device, dtype=weight_dtype)
    
    return ControlNetXS

ControlNetXS = get_ControlNetXS(base_model, cnxs_path, device, size_ratio=0.125, weight_dtype=weight_dtype)
vae = AutoencoderKL.from_pretrained(vae_path)
pipe = StableDiffusionXLInstantIDXSPipeline.from_pretrained(
    pretrained_model_name_or_path,
    vae=vae,
    unet=ControlNetXS,
    controlnet=None,
    torch_dtype=weight_dtype,
)

pipe.cuda(device=device, dtype=weight_dtype, use_xformers=True)
pipe.load_ip_adapter(image_proj_path, cross_attn_path)

pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.unet.config.ctrl_learn_time_embedding = True
pipe = pipe.to(args.device)




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 get_example():
    case = [
        [
            "./examples/1.jpg",
            None,
            "a woman,(looking at the viewer), portrait, daily wear, 8K texture, realistic, symmetrical hyperdetailed texture, masterpiece, enhanced details, (eye highlight:2),   perfect composition, natural lighting, best quality, authentic, natural posture",
            "(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",
        ],
        [
            "./examples/1.jpeg",
            "./examples/poses/pose1.jpg",
            "a woman,(looking at the viewer), portrait, daily wear, 8K texture, realistic, symmetrical hyperdetailed texture, masterpiece, enhanced details, (eye highlight:2),   perfect composition, natural lighting, best quality, authentic, natural posture",
            "(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",
        ],
    ]
    return case

def run_for_examples(face_file, pose_file, prompt, style, negative_prompt):
    return generate_image(
        face_file,
        pose_file,
        prompt,
        negative_prompt,
        20,  # num_steps
        0.8,  # identitynet_strength_ratio
        0.8,  # adapter_strength_ratio
        0.8,  # pose_strength
        5.0,  # guidance_scale
        42,  # seed
    )

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


@spaces.GPU
def generate_image(
    face_image_path,
    pose_image_path,
    prompt,
    negative_prompt,
    num_steps,
    controlnet_conditioning_scale,
    adapter_strength_ratio,
    guidance_scale,
    seed,
    progress=gr.Progress(track_tqdm=True),
):

    if face_image_path is None:
        raise gr.Error(
            f"Cannot find any input face image! Please upload the face image"
        )

    if prompt is None:
        prompt = "a person"

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

    face_image = load_image(face_image_path)
    face_image = resize_img(face_image, max_side=1024)
    face_image_cv2 = convert_from_image_to_cv2(face_image)
    height, width, _ = face_image_cv2.shape

    # Extract face features
    face_info = app.get(face_image_cv2)

    if len(face_info) == 0:
        raise gr.Error(
            f"Unable to detect a face in the image. Please upload a different photo with a clear face."
        )

    face_info = sorted(
        face_info,
        key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
    )[-1]  # only use the maximum face

    face_emb = torch.from_numpy(face_info.normed_embedding)
    face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
    img_controlnet = 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

    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=face_kps,
        control_mask=control_mask,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        height=height,
        width=width,
        generator=torch.Generator(device=device).manual_seed(seed),
    ).images

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



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

    with gr.Row():
        with gr.Column():
            with gr.Row(equal_height=True):
                # upload face image
                face_file = gr.Image(
                    label="Upload a photo of your face", type="filepath"
                )
                # optional: upload a reference pose image
                pose_file = gr.Image(
                    label="Upload a reference pose image (Optional)",
                    type="filepath",
                )

            # 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")
            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",
            )

            # strength
            controlnet_conditioning_scale = gr.Slider(
                label="IdentityNet strength (for fidelity)",
                minimum=0,
                maximum=1.0,
                step=0.1,
                value=0.8,
            )
            adapter_strength_ratio = gr.Slider(
                label="Image adapter strength (for detail)",
                minimum=0,
                maximum=1.2,
                step=0.1,
                value=0.8,
            )

            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=20,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5.0,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42,
                )

                randomize_seed = gr.Checkbox(label="Randomize seed", 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
            )

        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_file,
                pose_file,
                prompt,
                negative_prompt,
                num_steps,
                controlnet_conditioning_scale,
                adapter_strength_ratio,
                guidance_scale,
                seed,
            ],
            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_file, pose_file, prompt, negative_prompt],
        fn=run_for_examples,
        outputs=[gallery, usage_tips],
        cache_examples=True,
    )

    gr.Markdown(article)

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