initial demo

app.py

@@ -0,0 +1,231 @@
+import sys
+sys.path.append('./')
+
+from diffusers import (
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+    DPMSolverMultistepScheduler,
+)
+
+from arc2face import CLIPTextModelWrapper, project_face_embs
+
+import torch
+from insightface.app import FaceAnalysis
+from PIL import Image
+import numpy as np
+import random
+
+import gradio as gr
+
+# global variable
+MAX_SEED = np.iinfo(np.int32).max
+if torch.cuda.is_available():
+    device = "cuda"
+    dtype = torch.float16
+else:
+    device = "cpu"
+    dtype = torch.float32
+
+
+# download models
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="FoivosPar/Arc2Face", filename="arc2face/config.json", local_dir="./models")
+hf_hub_download(repo_id="FoivosPar/Arc2Face", filename="arc2face/diffusion_pytorch_model.safetensors", local_dir="./models")
+hf_hub_download(repo_id="FoivosPar/Arc2Face", filename="encoder/config.json", local_dir="./models")
+hf_hub_download(repo_id="FoivosPar/Arc2Face", filename="encoder/pytorch_model.bin", local_dir="./models")
+hf_hub_download(repo_id="FoivosPar/Arc2Face", filename="arcface.onnx", local_dir="./models/antelopev2")
+
+# Load face detection and recognition package
+if device=="cuda":
+    app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+else:
+    app = FaceAnalysis(name='antelopev2', root='./', providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+# Load pipeline
+base_model = 'runwayml/stable-diffusion-v1-5'
+encoder = CLIPTextModelWrapper.from_pretrained(
+    'models', subfolder="encoder", torch_dtype=dtype
+)
+unet = UNet2DConditionModel.from_pretrained(
+    'models', subfolder="arc2face", torch_dtype=dtype
+)
+pipeline = StableDiffusionPipeline.from_pretrained(
+        base_model,
+        text_encoder=encoder,
+        unet=unet,
+        torch_dtype=dtype,
+        safety_checker=None
+    )
+pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+pipeline = pipeline.to(device)
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+def get_example():
+    case = [
+        [
+            './assets/examples/freeman.jpg',
+        ],
+        [
+            './assets/examples/lily.png',
+        ],
+        [
+            './assets/examples/joacquin.png',
+        ],
+        [
+            './assets/examples/jackie.png',
+        ], 
+        [
+            './assets/examples/freddie.png',
+        ],
+        [
+            './assets/examples/hepburn.png',
+        ],
+    ]
+    return case
+
+def run_example(img_file):
+    return generate_image(img_file, 25, 3, 23, 2)
+
+
+def generate_image(image_path, num_steps, guidance_scale, seed, num_images, progress=gr.Progress(track_tqdm=True)):
+
+    if image_path is None:
+        raise gr.Error(f"Cannot find any input face image! Please upload a face image.")
+    
+    img = np.array(Image.open(image_path))[:,:,::-1]
+
+    # Face detection and ID-embedding extraction
+    faces = app.get(img)
+    
+    if len(faces) == 0:
+        raise gr.Error(f"Face detection failed! Please try with another image")
+    
+    faces = sorted(faces, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1]  # select largest face (if more than one detected)
+    id_emb = torch.tensor(faces['embedding'], dtype=dtype)[None].to(device)
+    id_emb = id_emb/torch.norm(id_emb, dim=1, keepdim=True)   # normalize embedding
+    id_emb = project_face_embs(pipeline, id_emb)    # pass throught the encoder
+                    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    print("Start inference...")        
+    images = pipeline(
+        prompt_embeds=id_emb,
+        num_inference_steps=num_steps,
+        guidance_scale=guidance_scale, 
+        num_images_per_prompt=num_images,
+        generator=generator
+    ).images
+
+    return images
+
+### Description
+title = r"""
+<h1>Arc2Face: A Foundation Model of Human Faces</h1>
+"""
+
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://arc2face.github.io/' target='_blank'><b>Arc2Face: A Foundation Model of Human Faces</b></a>.<br>
+
+Steps:<br>
+1. Upload an image with a face. If multiple faces are detected, we use the largest one. For images with already tightly cropped faces, detection may fail, try images with a larger margin.
+2. Click <b>Submit</b> to generate new images of the subject.
+"""
+
+Footer = r"""
+---
+📝 **Citation**
+<br>
+If you find Arc2Face helpful for your research, please consider citing our paper:
+```bibtex
+@misc{paraperas2024arc2face,
+      title={Arc2Face: A Foundation Model of Human Faces}, 
+      author={Foivos Paraperas Papantoniou and Alexandros Lattas and Stylianos Moschoglou and Jiankang Deng and Bernhard Kainz and Stefanos Zafeiriou},
+      year={2024},
+      eprint={2403.11641},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```
+"""
+
+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():
+            
+            # upload face image
+            img_file = gr.Image(label="Upload a photo with a face", type="filepath")
+            
+            submit = gr.Button("Submit", variant="primary")
+            
+            with gr.Accordion(open=False, label="Advanced Options"):
+                num_steps = gr.Slider( 
+                    label="Number of sample steps",
+                    minimum=20,
+                    maximum=100,
+                    step=1,
+                    value=25,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.1,
+                    maximum=10.0,
+                    step=0.1,
+                    value=3,
+                )
+                num_images = gr.Slider(
+                    label="Number of output images",
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    value=2,
+                )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        with gr.Column():
+            gallery = gr.Gallery(label="Generated Images")
+
+        submit.click(
+            fn=randomize_seed_fn,
+            inputs=[seed, randomize_seed],
+            outputs=seed,
+            queue=False,
+            api_name=False,
+        ).then(
+            fn=generate_image,
+            inputs=[img_file, num_steps, guidance_scale, seed, num_images],
+            outputs=[gallery]
+        )
+    
+    
+    gr.Examples(
+        examples=get_example(),
+        inputs=[img_file],
+        run_on_click=True,
+        fn=run_example,
+        outputs=[gallery],
+    )
+    
+    gr.Markdown(Footer)
+
+demo.launch()

arc2face/__init__.py

@@ -0,0 +1,2 @@
+from .models import CLIPTextModelWrapper
+from .utils import project_face_embs

arc2face/models.py

@@ -0,0 +1,91 @@
+import torch
+from transformers import CLIPTextModel
+from typing import Any, Callable, Dict, Optional, Tuple, Union, List
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.clip.modeling_clip import _make_causal_mask, _expand_mask
+
+
+class CLIPTextModelWrapper(CLIPTextModel):
+    # Adapted from https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/clip/modeling_clip.py#L812
+    # Modified to accept precomputed token embeddings "input_token_embs" as input or calculate them from input_ids and return them.
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        input_token_embs: Optional[torch.Tensor] = None,
+        return_token_embs: Optional[bool] = False,
+    ) -> Union[Tuple, torch.Tensor, BaseModelOutputWithPooling]:
+
+        if return_token_embs:
+            return self.text_model.embeddings.token_embedding(input_ids)
+    
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    
+        output_attentions = output_attentions if output_attentions is not None else self.text_model.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.text_model.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.text_model.config.use_return_dict
+    
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+    
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+    
+        hidden_states = self.text_model.embeddings(input_ids=input_ids, position_ids=position_ids, inputs_embeds=input_token_embs)
+    
+        # CLIP's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+        causal_attention_mask = _make_causal_mask(input_shape, hidden_states.dtype, device=hidden_states.device)
+        # expand attention_mask
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+    
+        encoder_outputs = self.text_model.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+    
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.text_model.final_layer_norm(last_hidden_state)
+    
+        if self.text_model.eos_token_id == 2:
+            # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+            # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+            # ------------------------------------------------------------
+            # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
+            ]
+        else:
+            # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+                (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.text_model.eos_token_id)
+                .int()
+                .argmax(dim=-1),
+            ]
+    
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+    
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

arc2face/utils.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn.functional as F
+
+@torch.no_grad()
+def project_face_embs(pipeline, face_embs):
+
+    '''
+    face_embs: (N, 512) normalized ArcFace embeddings
+    '''
+
+    arcface_token_id = pipeline.tokenizer.encode("id", add_special_tokens=False)[0]
+
+    input_ids = pipeline.tokenizer(
+            "photo of a id person",
+            truncation=True,
+            padding="max_length",
+            max_length=pipeline.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids.to(pipeline.device)
+
+    face_embs_padded = F.pad(face_embs, (0, pipeline.text_encoder.config.hidden_size-512), "constant", 0)
+    token_embs = pipeline.text_encoder(input_ids=input_ids.repeat(len(face_embs), 1), return_token_embs=True)
+    token_embs[input_ids==arcface_token_id] = face_embs_padded
+
+    prompt_embeds = pipeline.text_encoder(
+        input_ids=input_ids,
+        input_token_embs=token_embs
+    )[0]
+
+    return prompt_embeds

requirements.txt

@@ -0,0 +1,9 @@
+numpy<1.24.0
+torch==2.0.1
+torchvision==0.15.2
+diffusers==0.22.0
+transformers==4.34.1
+accelerate
+insightface
+onnxruntime-gpu
+gradio