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
import spaces
# global variable
MAX_SEED = np.iinfo(np.int32).max
if torch.cuda.is_available():
device = "cuda"
dtype = torch.float16
elif torch.backends.mps.is_available():
device = "mps"
#dtype = torch.float32
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
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)
@spaces.GPU
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"""
Arc2Face: A Foundation Model of Human Faces
"""
description = r"""
Official 🤗 Gradio demo for Arc2Face: A Foundation Model of Human Faces.
Steps:
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 Submit to generate new images of the subject.
"""
Footer = r"""
---
📝 **Citation**
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()