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_masks2025-03-06T03-31-21_zero3-ada-1000.pt')
parser.add_argument('--model_style_type', type=str, default='photorealistic',
                    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('--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

def is_running_on_spaces():
    return os.getenv("SPACE_ID") is not None

from huggingface_hub import snapshot_download
large_files = ["models/*", "models/**/*"]
snapshot_download(repo_id="adaface-neurips/adaface-animate-models", 
                  repo_type="model", allow_patterns=large_files, local_dir=".")
os.makedirs("/tmp/gradio", exist_ok=True)

# 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=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(320, 320))

if is_running_on_spaces():
    device = 'cuda:0'
else:
    if args.gpu is None:
        device = "cuda" 
    else:
        device = f"cuda:{args.gpu}"

print(f"Device: {device}")

global adaface, id_animator

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

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, highlight_face, guidance_scale, out_image_count=4):
    if uploaded_image_paths is None:
        print("No image uploaded")
        return None, None, None

    global adaface, id_animator
    adaface.to(device)
    id_animator.to(device)

    # 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]
    
    with torch.no_grad():
        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)

    if highlight_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

    guidance_scale = min(guidance_scale, 5)
    
    # samples: A list of PIL Image instances.
    with torch.no_grad():
        samples = adaface(noise, prompt, placeholder_tokens_pos='append',
                          guidance_scale=guidance_scale, 
                          out_image_count=out_image_count, 
                          repeat_prompt_for_each_encoder=True,
                          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, 
                   seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                   highlight_face, is_adaface_enabled, adaface_power_scale, 
                   id_animator_anneal_steps, progress=gr.Progress(track_tqdm=True)):
    
    global adaface, id_animator
    adaface.to(device)
    id_animator.to(device)

    if prompt is None:
        prompt = ""

    #prompt = prompt + " 8k uhd, high quality"
    #if " shot" not in prompt:
    #    prompt = prompt + ", medium shot"

    if highlight_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

    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
        # ID-Animator Image Embedding Initial and End Scales
        image_embed_cfg_scales = (1, 1)
    else:
        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',
                                      repeat_prompt_for_each_encoder=True,
                                      verbose=True)

        # ID-Animator Image Embedding Initial and End Scales
        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 >= 1.2, weird artifacts appear. 
                                  # Here it's limited to 1~1.1.
                                  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, progress=gr.Progress()):
    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='cpu')
        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='cpu')
        # 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 Human Subject-Driven Video Generation
        """
    )
    gr.Markdown(
        """
<b>Official demo</b> for our working paper <b>AdaFace: A Versatile Text-space Face Encoder for Face Synthesis and Processing</b>.<br>

❗️**NOTE**❗️
- Support switching between three model styles: **Realistic**, **Photorealistic** and **Anime**. **Realistic** is less realistic than **Photorealistic** but has better motions.
- If you change the model style, please wait for 20~30 seconds for loading new model weight before the model begins to generate images/videos.

❗️**Tips**❗️
- You can upload one or more subject images for generating ID-specific video.
- If the face loses focus, try enabling "Highlight face".
- If the motion is weird, e.g., the prompt is "... 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>

        """
    )

    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"
                    )
            files.GRADIO_CACHE = "/tmp/gradio"
            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 4 images and select 1 image",
                            file_types=["image"],
                            file_count="multiple"
                    )
            init_img_files.GRADIO_CACHE = "/tmp/gradio"
            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 4 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 'walking on the beach'.",
                    value="highlighted hair, futuristic silver armor suit, confident stance, living room, smiling, head tilted, perfect smooth skin",
                    allow_custom_value=True,
                    choices=[
                            "portrait, highlighted hair, futuristic silver armor suit, confident stance, living room, smiling, head tilted, perfect smooth skin",
                            "portrait, walking on the beach, sunset",
                            "portrait, in a white apron and chef hat, garnishing a gourmet dish",
                            "portrait, dancing pose among folks in a park, waving hands",
                            "portrait, in iron man costume, the sky ablaze with hues of orange and purple",
                            "portrait, jedi wielding a lightsaber, star wars",
                            "portrait, night view of tokyo street, neon light",
                            "portrait, playing guitar on a boat, ocean waves",
                            "portrait, with a passion for reading, curled up with a book in a cozy nook near a window",
                            "portrait, celebrating new year, fireworks",
                            "portrait, running pose in a park",
                            "portrait, in space suit, space helmet, walking on mars",
                            "portrait, in superman costume, the sky ablaze with hues of orange and purple"
                    ])

            highlight_face = gr.Checkbox(label="Highlight face", value=False, 
                                         info="Enhance the facial features by prepending 'face portrait' to the prompt",
                                         visible=True)
            
            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=3,
                    step=0.1,
                    value=1,
                )
            init_image_final_weight = gr.Slider(
                    label="Final Strength of the Init Image",
                    info="How much the init image should influence the end of the video",
                    minimum=0,
                    maximum=2,
                    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,
            )

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

            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_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.05,
                        value=1.1,
                        visible=True,
                    )

                attn_scale = gr.Slider(
                        label="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.5,
                        maximum=2,
                        step=0.1,
                        value=1,
                        visible=True
                    )

                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,
                        maximum=1,
                        step=0.1,
                        value=0.5,
                    )
                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,
                        maximum=1,
                        step=0.1,
                        value=0.1,
                    )            

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

                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, nude, naked, nsfw, topless, bare breasts",
                )

        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=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=gen_init_images, inputs=[uploaded_files_gallery, prompt, highlight_face,
                                                   guidance_scale], 
                       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,
                         seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                         highlight_face, is_adaface_enabled, 
                         adaface_power_scale, id_animator_anneal_steps],
                 outputs=[result_video]
        )

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