Cothn/PhotoMaker
1
1import torch2import numpy as np3import random4import os5 6from diffusers.utils import load_image7from diffusers import DDIMScheduler8 9from huggingface_hub import hf_hub_download10import spaces11import gradio as gr12 13from pipeline import PhotoMakerStableDiffusionXLPipeline14from style_template import styles15 16# global variable17base_model_path = 'SG161222/RealVisXL_V3.0'18device = "cuda" if torch.cuda.is_available() else "cpu"19MAX_SEED = np.iinfo(np.int32).max20STYLE_NAMES = list(styles.keys())21DEFAULT_STYLE_NAME = "Photographic (Default)"22 23# download PhotoMaker checkpoint to cache24photomaker_ckpt = hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")25 26pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(27 base_model_path, 28 torch_dtype=torch.bfloat16, 29 use_safetensors=True, 30 variant="fp16",31).to(device)32 33pipe.load_photomaker_adapter(34 os.path.dirname(photomaker_ckpt),35 subfolder="",36 weight_name=os.path.basename(photomaker_ckpt),37 trigger_word="img"38) 39pipe.id_encoder.to(device)40 41pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)42# pipe.set_adapters(["photomaker"], adapter_weights=[1.0])43pipe.fuse_lora()44 45@spaces.GPU46def generate_image(upload_images, prompt, negative_prompt, style_name, num_steps, style_strength_ratio, num_outputs, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):47 # check the trigger word48 image_token_id = pipe.tokenizer.convert_tokens_to_ids(pipe.trigger_word)49 input_ids = pipe.tokenizer.encode(prompt)50 if image_token_id not in input_ids:51 raise gr.Error(f"Cannot find the trigger word '{pipe.trigger_word}' in text prompt! Please refer to step 2️⃣")52 53 if input_ids.count(image_token_id) > 1:54 raise gr.Error(f"Cannot use multiple trigger words '{pipe.trigger_word}' in text prompt!")55 56 # apply the style template57 prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)58 59 if upload_images is None:60 raise gr.Error(f"Cannot find any input face image! Please refer to step 1️⃣")61 62 input_id_images = []63 for img in upload_images:64 input_id_images.append(load_image(img))65 66 generator = torch.Generator(device=device).manual_seed(seed)67 68 print("Start inference...")69 print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")70 start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)71 if start_merge_step > 30:72 start_merge_step = 3073 print(start_merge_step)74 images = pipe(75 prompt=prompt,76 input_id_images=input_id_images,77 negative_prompt=negative_prompt,78 num_images_per_prompt=num_outputs,79 num_inference_steps=num_steps,80 start_merge_step=start_merge_step,81 generator=generator,82 guidance_scale=guidance_scale,83 ).images84 return images, gr.update(visible=True)85 86def swap_to_gallery(images):87 return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)88 89def upload_example_to_gallery(images, prompt, style, negative_prompt):90 return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)91 92def remove_back_to_files():93 return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)94 95def remove_tips():96 return gr.update(visible=False)97 98def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:99 if randomize_seed:100 seed = random.randint(0, MAX_SEED)101 return seed102 103def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:104 p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])105 return p.replace("{prompt}", positive), n + ' ' + negative106 107def get_image_path_list(folder_name):108 image_basename_list = os.listdir(folder_name)109 image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])110 return image_path_list111 112def get_example():113 case = [114 [115 get_image_path_list('./examples/scarletthead_woman'),116 "instagram photo, portrait photo of a woman img, colorful, perfect face, natural skin, hard shadows, film grain",117 "(No style)",118 "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",119 ],120 [121 get_image_path_list('./examples/newton_man'),122 "sci-fi, closeup portrait photo of a man img wearing the sunglasses in Iron man suit, face, slim body, high quality, film grain",123 "(No style)",124 "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",125 ],126 ]127 return case128 129### Description and style130logo = r"""131<center><img src='https://photo-maker.github.io/assets/logo.png' alt='PhotoMaker logo' style="width:80px; margin-bottom:10px"></center>132"""133title = r"""134<h1 align="center">PhotoMaker: Customizing Realistic Human Photos via Stacked ID Embedding</h1>135"""136 137description = r"""138<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/TencentARC/PhotoMaker' target='_blank'><b>PhotoMaker: Customizing Realistic Human Photos via Stacked ID Embedding</b></a>.<br>139<br>140For stylization, you could use our other gradio demo [PhotoMaker-Style](https://huggingface.co/spaces/TencentARC/PhotoMaker-Style).141<br>142❗️❗️❗️[<b>Important</b>] Personalization steps:<br>1431️⃣ Upload images of someone you want to customize. One image is ok, but more is better. Although we do not perform face detection, the face in the uploaded image should <b>occupy the majority of the image</b>.<br>1442️⃣ Enter a text prompt, making sure to <b>follow the class word</b> you want to customize with the <b>trigger word</b>: `img`, such as: `man img` or `woman img` or `girl img`.<br>1453️⃣ Choose your preferred style template.<br>1464️⃣ Click the <b>Submit</b> button to start customizing.147"""148 149article = r"""150 151If PhotoMaker is helpful, please help to ⭐ the <a href='https://github.com/TencentARC/PhotoMaker' target='_blank'>Github Repo</a>. Thanks! 152[](https://github.com/TencentARC/PhotoMaker)153---154📝 **Citation**155<br>156If our work is useful for your research, please consider citing:157 158```bibtex159@article{li2023photomaker,160 title={PhotoMaker: Customizing Realistic Human Photos via Stacked ID Embedding},161 author={Li, Zhen and Cao, Mingdeng and Wang, Xintao and Qi, Zhongang and Cheng, Ming-Ming and Shan, Ying},162 booktitle={arXiv preprint arxiv:2312.04461},163 year={2023}164}165```166📋 **License**167<br>168Apache-2.0 LICENSE. Please refer to the [LICENSE file](https://huggingface.co/TencentARC/PhotoMaker/blob/main/LICENSE) for details.169 170📧 **Contact**171<br>172If you have any questions, please feel free to reach me out at <b>zhenli1031@gmail.com</b>.173"""174 175tips = r"""176### Usage tips of PhotoMaker1771. Upload more photos of the person to be customized to **improve ID fidelty**. If the input is Asian face(s), maybe consider adding 'asian' before the class word, e.g., `asian woman img`1782. When stylizing, does the generated face look too realistic? Try switching to our **other gradio demo** [PhotoMaker-Style](https://huggingface.co/spaces/TencentARC/PhotoMaker-Style). Adjust the **Style strength** to 30-50, the larger the number, the less ID fidelty, but the stylization ability will be better.1793. For **faster** speed, reduce the number of generated images and sampling steps. However, please note that reducing the sampling steps may compromise the ID fidelity.180"""181# We have provided some generate examples and comparisons at: [this website]().182# 3. Don't make the prompt too long, as we will trim it if it exceeds 77 tokens. 183# 4. When generating realistic photos, if it's not real enough, try switching to our other gradio application [PhotoMaker-Realistic]().184 185css = '''186.gradio-container {width: 85% !important}187'''188with gr.Blocks(css=css) as demo:189 gr.Markdown(logo)190 gr.Markdown(title)191 gr.Markdown(description)192 # gr.DuplicateButton(193 # value="Duplicate Space for private use ",194 # elem_id="duplicate-button",195 # visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",196 # )197 with gr.Row():198 with gr.Column():199 files = gr.Files(200 label="Drag (Select) 1 or more photos of your face",201 file_types=["image"]202 )203 uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)204 with gr.Column(visible=False) as clear_button:205 remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")206 prompt = gr.Textbox(label="Prompt",207 info="Try something like 'a photo of a man/woman img', 'img' is the trigger word.",208 placeholder="A photo of a [man/woman img]...")209 style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)210 submit = gr.Button("Submit")211 212 with gr.Accordion(open=False, label="Advanced Options"):213 negative_prompt = gr.Textbox(214 label="Negative Prompt", 215 placeholder="low quality",216 value="nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry",217 )218 num_steps = gr.Slider( 219 label="Number of sample steps",220 minimum=20,221 maximum=100,222 step=1,223 value=50,224 )225 style_strength_ratio = gr.Slider(226 label="Style strength (%)",227 minimum=15,228 maximum=50,229 step=1,230 value=20,231 )232 num_outputs = gr.Slider(233 label="Number of output images",234 minimum=1,235 maximum=4,236 step=1,237 value=2,238 )239 guidance_scale = gr.Slider(240 label="Guidance scale",241 minimum=0.1,242 maximum=10.0,243 step=0.1,244 value=5,245 )246 seed = gr.Slider(247 label="Seed",248 minimum=0,249 maximum=MAX_SEED,250 step=1,251 value=0,252 )253 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)254 with gr.Column():255 gallery = gr.Gallery(label="Generated Images")256 usage_tips = gr.Markdown(label="Usage tips of PhotoMaker", value=tips ,visible=False)257 258 files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])259 remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])260 261 submit.click(262 fn=remove_tips,263 outputs=usage_tips, 264 ).then(265 fn=randomize_seed_fn,266 inputs=[seed, randomize_seed],267 outputs=seed,268 queue=False,269 api_name=False,270 ).then(271 fn=generate_image,272 inputs=[files, prompt, negative_prompt, style, num_steps, style_strength_ratio, num_outputs, guidance_scale, seed],273 outputs=[gallery, usage_tips]274 )275 276 gr.Examples(277 examples=get_example(),278 inputs=[files, prompt, style, negative_prompt],279 run_on_click=True,280 fn=upload_example_to_gallery,281 outputs=[uploaded_files, clear_button, files],282 )283 284 gr.Markdown(article)285 286demo.launch()