Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
import numpy as np | |
import random | |
import spaces # Uncomment if using ZeroGPU | |
import os | |
from diffusers import StableDiffusionPipeline, DDPMScheduler | |
import torch | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
model_repo_id = "stabilityai/stable-diffusion-2-1-base" | |
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype) | |
# pipe = pipe.to(device) | |
pipe = StableDiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.float16).to(device) | |
pipe.scheduler = DDPMScheduler.from_pretrained(model_repo_id, subfolder="scheduler") | |
folder_of_lora_weights = "./ID-Booth_LoRA_weights" | |
which_checkpoint = "checkpoint-31-6400" | |
lora_name = "pytorch_lora_weights.safetensors" | |
folder_of_identity_images = "./assets/example_images/" | |
backgrounds_list = ["Forest", "City street", "Beach", "Office", "Bus", "Laboratory", "Factory", "Construction site", "Hospital", "Night club", ""] | |
poses_list = ["Portrait", "Side-portrait"] | |
id_list = ["ID_1", "ID_5", "ID_16", "ID_20"] | |
gender_dict = {"ID_1": "male", "ID_5": "male", "ID_16": "female", "ID_20": "male"} | |
MAX_SEED = 10000 | |
image_size = 512 | |
# Uncomment if using ZeroGPU | |
def infer( | |
identity, | |
background, | |
pose, | |
negative_prompt, | |
seed, | |
randomize_seed, | |
guidance_scale, | |
num_inference_steps, | |
num_images=1 | |
): | |
full_lora_weights_path = f"{folder_of_lora_weights}/{identity}/{which_checkpoint}/{lora_name}" | |
pipe.load_lora_weights(full_lora_weights_path) | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
generator = torch.Generator().manual_seed(seed) | |
gender = gender_dict[identity] | |
# Construct prompt from dropdown selections | |
prompt = f"face {pose.lower()} photo of {gender} sks person, {background.lower()} background" | |
images = pipe( | |
prompt=prompt, | |
negative_prompt=negative_prompt, | |
guidance_scale=guidance_scale, | |
num_inference_steps=num_inference_steps, | |
width=image_size, | |
height=image_size, | |
generator=generator, | |
num_images_per_prompt=num_images, | |
).images | |
return images | |
### Description | |
header = " # ID-Booth: Identity-consistent Face Generation with Diffusion Models" | |
description = "This is an official Gradio demo for the paper <a href='https://dariant.github.io/publications/ID-Booth' target='_blank'>ID-Booth: Identity-consistent Face Generation with Diffusion Models</a>" | |
footer = r""" | |
**Citation** | |
<br> | |
If you find ID-Booth helpful, please consider citing our paper: | |
```bibtex | |
@article{tomasevic2025IDBooth, | |
title={{ID-Booth}: Identity-consistent Face Generation with Diffusion Models}, | |
author={Toma{\v{s}}evi{\'c}, Darian and Boutros, Fadi and Lin, Chenhao and Damer, Naser and {\v{S}}truc, Vitomir and Peer, Peter}, | |
journal={arXiv preprint arXiv:2504.07392}, | |
year={2025} | |
} | |
``` | |
""" | |
css = ''' | |
.gradio-container { | |
width: 75%; | |
margin: auto; | |
} | |
''' | |
with gr.Blocks(css=css) as demo: | |
# description | |
gr.Markdown(header) | |
gr.Markdown(description) | |
with gr.Column(): | |
# with gr.Row(): | |
# gr.Markdown("### Choose an identity, background, and pose:") | |
with gr.Row(): | |
for id in id_list: | |
image_path = os.path.join(folder_of_identity_images, id + ".jpg") | |
img = gr.Image(value=image_path, label=id, | |
width=256, height=256, | |
show_label=True, interactive=False, | |
show_download_button=False, | |
show_fullscreen_button=False, | |
show_share_button=False, | |
) | |
with gr.Row(): | |
identity = gr.Dropdown( | |
label="Identity:", | |
choices=id_list, | |
value=id_list[2], | |
) | |
background = gr.Dropdown( | |
label="Background:", | |
choices=backgrounds_list, | |
value=backgrounds_list[1], | |
) | |
pose = gr.Dropdown( | |
label="Pose:", | |
choices=poses_list, | |
value=poses_list[0], | |
) | |
run_button = gr.Button("Generate in-the-wild images", scale=0, variant="primary") | |
#result = gr.Image(label="Result", show_label=False) | |
result = gr.Gallery(label="Generated Images", show_label=False) | |
with gr.Accordion(open=False, label="Advanced Options"): | |
negative_prompt = gr.Text( | |
label="Negative prompt", | |
max_lines=1, | |
placeholder="Enter a negative prompt", | |
value="cartoon, cgi, render, illustration, painting, drawing, black and white, bad body proportions, landscape", | |
) | |
num_inference_steps = gr.Slider( | |
label="Number of sampling steps", | |
minimum=1, | |
maximum=100, | |
step=1, | |
value=30, | |
) | |
guidance_scale = gr.Slider( | |
label="Guidance scale", | |
minimum=0.1, | |
maximum=10.0, | |
step=0.1, | |
value=5.0, | |
) | |
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) | |
# gr.Examples( | |
# examples=[ | |
# [id_list[0], backgrounds_list[0], poses_list[0], "A beautiful photo of a person", 0, False, 512, 512, 7.5, 50], | |
# ], | |
# inputs=[selected_identity, background, pose], | |
# ) | |
gr.on( | |
triggers=[run_button.click], | |
fn=infer, | |
inputs=[ | |
identity, | |
background, | |
pose, | |
negative_prompt, | |
seed, | |
randomize_seed, | |
guidance_scale, | |
num_inference_steps, | |
num_images | |
], | |
outputs=[result], | |
) | |
gr.Markdown(footer) | |
if __name__ == "__main__": | |
demo.launch() | |