app.py

import gradio as gr
import spaces
css = '''
.gradio-container {width: 85% !important}
'''
from animatediff.utils.util import save_videos_grid
from adaface.adaface_wrapper import AdaFaceWrapper

import random
from infer import load_model, model_style_type2base_model_path
MAX_SEED=10000
import uuid
from insightface.app import FaceAnalysis
import os
import os
import cv2
from diffusers.utils import load_image
from insightface.utils import face_align
from PIL import Image
import torch
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--adaface_encoder_types", type=str, nargs="+", default=["consistentID", "arc2face"],
                    choices=["arc2face", "consistentID"], help="Type(s) of the ID2Ada prompt encoders")
parser.add_argument('--adaface_ckpt_path', type=str, 
                    default='models/adaface/VGGface2_HQ_masks2024-10-08T14-42-05_zero3-ada-30000.pt')
parser.add_argument('--model_style_type', type=str, default='realistic',
                    choices=["realistic", "anime", "photorealistic"], help="Type of the base model")
parser.add_argument("--guidance_scale", type=float, default=8.0,
                    help="The guidance scale for the diffusion model. Default: 8.0")
parser.add_argument("--do_neg_id_prompt_weight", type=float, default=0,
                    help="The weight of added ID prompt embeddings into the negative prompt. Default: 0, disabled.")

parser.add_argument('--gpu', type=int, default=None)
parser.add_argument('--ip', type=str, default="0.0.0.0")
args = parser.parse_args()

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

# model = load_model()
# This FaceAnalysis is just to crop the face areas from the uploaded images,
# and is independent of the adaface FaceAnalysis apps.
app = FaceAnalysis(name="buffalo_l", root='models/insightface', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(320, 320))
device = "cuda" if args.gpu is None else f"cuda:{args.gpu}"

global adaface, id_animator

base_model_path = model_style_type2base_model_path[args.model_style_type]
id_animator = load_model(model_style_type=args.model_style_type, device=device)
adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
                         adaface_encoder_types=args.adaface_encoder_types,
                         adaface_ckpt_paths=[args.adaface_ckpt_path], device=device)

basedir = os.getcwd()
savedir = os.path.join(basedir,'samples')
os.makedirs(savedir, exist_ok=True)

#print(f"### Cleaning cached examples ...")
#os.system(f"rm -rf gradio_cached_examples/")

def swap_to_gallery(images):
    # Update uploaded_files_gallery, show files, hide clear_button_column
    # Or:
    # Update uploaded_init_img_gallery, show init_img_files, hide init_clear_button_column
    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(value=images, visible=False)

def remove_back_to_files():
    # Hide uploaded_files_gallery,    show clear_button_column,      hide files,           reset init_img_selected_idx
    # Or:
    # Hide uploaded_init_img_gallery, hide init_clear_button_column, show init_img_files,  reset init_img_selected_idx
    return gr.update(visible=False), gr.update(visible=False), gr.update(value=None, visible=True), gr.update(value="0")

def get_clicked_image(data: gr.SelectData):
    return data.index
    
@spaces.GPU
def gen_init_images(uploaded_image_paths, prompt, 
                    guidance_scale, do_neg_id_prompt_weight, out_image_count=3):
    global adaface, id_animator
    if uploaded_image_paths is None:
        print("No image uploaded")
        return None, None, None

    # uploaded_image_paths is a list of tuples:
    # [('/tmp/gradio/249981e66a7c665aaaf1c7eaeb24949af4366c88/jensen huang.jpg', None)]
    # Extract the file paths.
    uploaded_image_paths = [path[0] for path in uploaded_image_paths]
    adaface_subj_embs = \
        adaface.prepare_adaface_embeddings(image_paths=uploaded_image_paths, face_id_embs=None, 
                                           update_text_encoder=True)
    
    if adaface_subj_embs is None:
        raise gr.Error(f"Failed to detect any faces! Please try with other images")
        
    # Generate two images each time for the user to select from.
    noise = torch.randn(out_image_count, 3, 512, 512)

    enhance_face = True
    if enhance_face and "face portrait" not in prompt:
        if "portrait" in prompt:
            # Enhance the face features by replacing "portrait" with "face portrait".
            prompt = prompt.replace("portrait", "face portrait")
        else:
            prompt = "face portrait, " + prompt

    # samples: A list of PIL Image instances.
    with torch.no_grad():
        samples = adaface(noise, prompt, placeholder_tokens_pos='append',
                          guidance_scale=guidance_scale, 
                          do_neg_id_prompt_weight=do_neg_id_prompt_weight,
                          out_image_count=out_image_count, verbose=True)

    face_paths = []
    for sample in samples:        
        random_name = str(uuid.uuid4())
        face_path = os.path.join(savedir, f"{random_name}.jpg")
        face_paths.append(face_path)
        sample.save(face_path)
        print(f"Generated init image: {face_path}")

    # Update uploaded_init_img_gallery, update and hide init_img_files, hide init_clear_button_column
    return gr.update(value=face_paths, visible=True), gr.update(value=face_paths, visible=False), gr.update(visible=True)

@spaces.GPU(duration=90)
def generate_video(image_container, uploaded_image_paths, init_img_file_paths, init_img_selected_idx, 
                   init_image_strength, init_image_final_weight,
                   prompt, negative_prompt, num_steps, video_length, guidance_scale, do_neg_id_prompt_weight, 
                   seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                   is_adaface_enabled, adaface_ckpt_path, adaface_power_scale, 
                   id_animator_anneal_steps, progress=gr.Progress(track_tqdm=True)):
    global adaface, id_animator

    if prompt is None:
        prompt = ""

    prompt = prompt + " 8k uhd, high quality"
    if " shot" not in prompt:
        prompt = prompt + ", medium shot"
        
    prompt_img_lists=[]
    for path in uploaded_image_paths:
        img = cv2.imread(path)
        faces = app.get(img)
        face_roi = face_align.norm_crop(img, faces[0]['kps'], 112)
        random_name = str(uuid.uuid4())
        face_path = os.path.join(savedir, f"{random_name}.jpg")
        cv2.imwrite(face_path, face_roi)
        # prompt_img_lists is a list of PIL images.
        prompt_img_lists.append(load_image(face_path).resize((224,224)))

    if adaface is None or not is_adaface_enabled:
        adaface_prompt_embeds, negative_prompt_embeds = None, None
        image_embed_cfg_scales = (1, 1)
    else:
        if (adaface_ckpt_path is not None and adaface_ckpt_path.strip() != '') \
          and (adaface_ckpt_path != args.adaface_ckpt_path):
            args.adaface_ckpt_path = adaface_ckpt_path
            # Reload the adaface model weights.
            adaface.id2ada_prompt_encoder.load_adaface_ckpt(adaface_ckpt_path)

        with torch.no_grad():
            adaface_subj_embs = \
                adaface.prepare_adaface_embeddings(image_paths=uploaded_image_paths, face_id_embs=None, 
                                                   update_text_encoder=True)

            # adaface_prompt_embeds: [1, 77, 768].
            adaface_prompt_embeds, negative_prompt_embeds, _, _ = \
                adaface.encode_prompt(prompt, placeholder_tokens_pos='append',
                                      do_neg_id_prompt_weight=do_neg_id_prompt_weight,
                                      verbose=True)

        image_embed_cfg_scales = (image_embed_cfg_begin_scale, image_embed_cfg_end_scale)

    # init_img_file_paths is a list of image paths. If not chose, init_img_file_paths is None.
    if init_img_file_paths is not None:
        init_img_selected_idx = int(init_img_selected_idx)
        init_img_file_path = init_img_file_paths[init_img_selected_idx]
        init_image = cv2.imread(init_img_file_path)
        init_image = cv2.resize(init_image, (512, 512))
        init_image = Image.fromarray(cv2.cvtColor(init_image, cv2.COLOR_BGR2RGB))
        print(f"init_image: {init_img_file_path}")
    else:
        init_image = None

    sample = id_animator.generate(prompt_img_lists, 
                                  init_image            = init_image,
                                  init_image_strength   = (init_image_strength, init_image_final_weight),
                                  prompt                = prompt,
                                  negative_prompt       = negative_prompt,
                                  adaface_prompt_embeds = (adaface_prompt_embeds, negative_prompt_embeds),
                                  # adaface_power_scale is not so useful, and when it's set >= 2, weird artifacts appear. 
                                  # Here it's limited to 0.7~1.3.
                                  adaface_power_scale   = adaface_power_scale,
                                  num_inference_steps   = num_steps,
                                  id_animator_anneal_steps  = id_animator_anneal_steps,
                                  seed                  = seed,
                                  guidance_scale        = guidance_scale,
                                  width                 = 512,
                                  height                = 512,
                                  video_length          = video_length,
                                  attn_scale            = attn_scale,
                                  image_embed_cfg_scales = image_embed_cfg_scales,
                                )
    
    save_sample_path = os.path.join(savedir, f"{random_name}.mp4")
    save_videos_grid(sample, save_sample_path)
    return save_sample_path

def check_prompt_and_model_type(prompt, model_style_type):
    global adaface, id_animator

    model_style_type = model_style_type.lower()
    base_model_path = model_style_type2base_model_path[model_style_type]
    # If the base model type is changed, reload the model.
    if model_style_type != args.model_style_type:
        id_animator = load_model(model_style_type=model_style_type, device=device)
        adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
                                 adaface_encoder_types=args.adaface_encoder_types,
                                 adaface_ckpt_paths=[args.adaface_ckpt_path], device=device)
        # Update base model type.
        args.model_style_type = model_style_type

    if not prompt:
        raise gr.Error("Prompt cannot be blank")

with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
    gr.Markdown(
        """
        # AdaFace-Animate: Zero-Shot Subject-Driven Video Generation for Humans
        """
    )
    gr.Markdown(
        """
<b>Official demo</b> for our working paper <b>AdaFace: A Versatile Face Encoder for Zero-Shot Diffusion Model Personalization</b>.<br>

❗️**Tips**❗️
- You can upload one or more subject images for generating ID-specific video.
- If the face dominates the video frames, try increasing the Weight of ID prompt in the negative prompt, at the cost of slight drop of ID authenticity.
- If the face loses focu, try increasing the guidance scale. At the same time, increase the Weight of ID prompt in the negative prompt proportionally.
- If the motion is weird, e.g., running, try increasing the number of sampling steps.
- Usage explanations and demos: [Readme](https://huggingface.co/spaces/adaface-neurips/adaface-animate/blob/main/README2.md).
- AdaFace Text-to-Image: <a href="https://huggingface.co/spaces/adaface-neurips/adaface" style="display: inline-flex; align-items: center;">
  AdaFace 
  <img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-yellow" alt="Hugging Face Spaces" style="margin-left: 5px;">
  </a>

**TODO:**
- ControlNet integration.
        """
    )

    with gr.Row():
        with gr.Column():
            files = gr.File(
                        label="Drag / Select 1 or more photos of a person's face",
                        file_types=["image"],
                        file_count="multiple"
                    )
            image_container = gr.Image(label="image container", sources="upload", type="numpy", height=256, visible=False)
            uploaded_files_gallery = gr.Gallery(label="Subject images", visible=False, columns=3, rows=2, height=300)
            with gr.Column(visible=False) as clear_button_column:
                remove_and_reupload = gr.ClearButton(value="Remove and upload subject images", components=files, size="sm")

            init_img_files = gr.File(
                            label="[Optional] Generate 3 images and select 1 image",
                            file_types=["image"],
                            file_count="multiple"
                    )
            init_img_container = gr.Image(label="init image container", sources="upload", type="numpy", height=256, visible=False)
            # Although there's only one image, we still use columns=3, to scale down the image size.
            # Otherwise it will occupy the full width, and the gallery won't show the whole image.
            uploaded_init_img_gallery = gr.Gallery(label="Init image", visible=False, columns=3, rows=1, height=200)
            # placeholder is just hint, not the real value. So we use "value='0'" instead of "placeholder='0'".
            init_img_selected_idx = gr.Textbox(label="Selected init image index", value="0", visible=False)

            with gr.Column(visible=True) as init_gen_button_column:
                gen_init = gr.Button(value="Generate 3 new init images")
            with gr.Column(visible=False) as init_clear_button_column:
                remove_init_and_reupload = gr.ClearButton(value="Upload an old init image", components=init_img_files, size="sm")

            prompt = gr.Dropdown(label="Prompt",
                       info="Try something like 'man/woman walking on the beach'.",
                       value="((best quality)), ((masterpiece)), ((realistic)), highlighted hair, futuristic silver armor suit, confident stance, high-resolution, living room, smiling, head tilted, perfect smooth skin",
                       allow_custom_value=True,
                       filterable=False,
                       choices=[
                            "((best quality)), ((masterpiece)), ((realistic)), highlighted hair, futuristic silver armor suit, confident stance, high-resolution, living room, smiling, head tilted, perfect smooth skin",
                            "walking on the beach, sunset, orange sky, eye level shot",
                            "in a white apron and chef hat, garnishing a gourmet dish, full body view, long shot",
                            "dancing pose among folks in a park, waving hands",
                            "in iron man costume flying pose, the sky ablaze with hues of orange and purple, full body view, long shot",
                            "jedi wielding a lightsaber, star wars, full body view, eye level shot",
                            "playing guitar on a boat, ocean waves",
                            "with a passion for reading, curled up with a book in a cozy nook near a window",
                            #"running pose in a park, full body view, eye level shot",
                            "in superman costume flying pose, the sky ablaze with hues of orange and purple, full body view, long shot"
                       ])

            init_image_strength = gr.Slider(
                    label="Init Image Strength",
                    info="How much the init image should influence each frame. 0: no influence (scenes are more dynamic), 3: strongest influence (scenes are more static).",
                    minimum=0,
                    maximum=1.5,
                    step=0.25,
                    value=1,
                )
            init_image_final_weight = gr.Slider(
                    label="Final Weight of the Init Image",
                    info="How much the init image should influence the end of the video",
                    minimum=0,
                    maximum=0.25,
                    step=0.025,
                    value=0.1,
                )

            model_style_type = gr.Dropdown(
                label="Base Model Style Type",
                info="Switching the base model type will take 10~20 seconds to reload the model",
                value=args.model_style_type.capitalize(),
                choices=["Realistic", "Anime"], #"Photorealistic"],
                allow_custom_value=False,
                filterable=False,
            )
            guidance_scale = gr.Slider(
                label="Guidance scale",
                info="If > 10, there may be artifacts.",
                minimum=1.0,
                maximum=12.0,
                step=1,
                value=args.guidance_scale,
            )

            do_neg_id_prompt_weight = gr.Slider(
                label="Weight of ID prompt in the negative prompt",
                minimum=0.0,
                maximum=0.9,
                step=0.1,
                value=args.do_neg_id_prompt_weight,
                visible=False
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=985,
            )
            randomize_seed = gr.Checkbox(
                label="Randomize seed", 
                value=True, 
                info="Uncheck for reproducible results")

            negative_prompt = gr.Textbox(
                label="Negative Prompt", 
                placeholder="low quality",
                value="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime), text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, bare breasts, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, long neck, UnrealisticDream",
            )
            num_steps = gr.Slider( 
                label="Number of sampling steps. More steps for better composition, but longer time.",
                minimum=30,
                maximum=70,
                step=10,
                value=40,
            )

            submit = gr.Button("Generate Video")

            with gr.Accordion(open=False, label="Advanced Options"):
                video_length = gr.Slider(
                    label="video_length",
                    info="Do not change; any values other than 16 will mess up the output video",
                    minimum=16,
                    maximum=21,
                    step=1,
                    value=16,
                    interactive=False,
                    visible=False,
                )
                is_adaface_enabled = gr.Checkbox(label="Enable AdaFace", 
                                                 info="Enable AdaFace for better face details. If unchecked, it falls back to ID-Animator (https://huggingface.co/spaces/ID-Animator/ID-Animator).",
                                                 value=True)
                adaface_ckpt_path = gr.Textbox(
                    label="AdaFace checkpoint path", 
                    placeholder=args.adaface_ckpt_path,
                    value=args.adaface_ckpt_path,
                )

                adaface_power_scale = gr.Slider(
                        label="AdaFace Embedding Power Scale",
                        info="Increase this scale slightly only if the face is defocused or the face details are not clear",
                        minimum=0.8,
                        maximum=1.2,
                        step=0.1,
                        value=1,
                    )

                image_embed_cfg_begin_scale = gr.Slider(
                        label="ID-Animator Image Embedding Initial Scale",
                        info="The scale of the ID-Animator image embedding (influencing coarse facial features and poses)",
                        minimum=0.6,
                        maximum=1.5,
                        step=0.1,
                        value=1.0,
                    )
                image_embed_cfg_end_scale = gr.Slider(
                        label="ID-Animator Image Embedding Final Scale",
                        info="The scale of the ID-Animator image embedding (influencing coarse facial features and poses)",
                        minimum=0.3,
                        maximum=1.5,
                        step=0.1,
                        value=0.5,
                    )            

                id_animator_anneal_steps = gr.Slider(
                    label="ID-Animator Scale Anneal Steps",
                    minimum=0,
                    maximum=40,
                    step=1,
                    value=20,
                    visible=True,
                )

                attn_scale = gr.Slider(
                        label="ID-Animator Attention Processor Scale",
                        info="The scale of the ID embeddings on the attention (the higher, the more focus on the face, less on the background)" ,
                        minimum=0,
                        maximum=2,
                        step=0.1,
                        value=1,
                    )

        with gr.Column():
            result_video = gr.Video(label="Generated Animation", interactive=False)
        
        files.upload(fn=swap_to_gallery, inputs=files,     outputs=[uploaded_files_gallery, clear_button_column, files])
        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files_gallery, clear_button_column, files, init_img_selected_idx])

        init_img_files.upload(fn=swap_to_gallery, inputs=init_img_files, 
                              outputs=[uploaded_init_img_gallery, init_clear_button_column, init_img_files])
        remove_init_and_reupload.click(fn=remove_back_to_files, 
                                       outputs=[uploaded_init_img_gallery, init_clear_button_column, 
                                                init_img_files, init_img_selected_idx])
        gen_init.click(fn=gen_init_images, inputs=[uploaded_files_gallery, prompt, 
                                                   guidance_scale, do_neg_id_prompt_weight], 
                       outputs=[uploaded_init_img_gallery, init_img_files, init_clear_button_column])
        uploaded_init_img_gallery.select(fn=get_clicked_image, inputs=None, outputs=init_img_selected_idx)

        submit.click(fn=check_prompt_and_model_type,
                     inputs=[prompt, model_style_type],outputs=None).success(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
        ).then(
                 fn=generate_video,
                 inputs=[image_container, files, 
                         init_img_files, init_img_selected_idx, init_image_strength, init_image_final_weight,
                         prompt, negative_prompt, num_steps, video_length, guidance_scale, do_neg_id_prompt_weight,
                         seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                         is_adaface_enabled, adaface_ckpt_path, adaface_power_scale, id_animator_anneal_steps],
                 outputs=[result_video]
        )

demo.launch(share=True, server_name=args.ip, ssl_verify=False)