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 @spaces.GPU # 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 ID-Booth: Identity-consistent Face Generation with Diffusion Models" footer = r""" **Citation**
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()