Fantasy-Studio/Paint-by-Example
95
1import gradio as gr2 3from io import BytesIO4import requests5import PIL6from PIL import Image7import numpy as np8import os9import uuid10import torch11from torch import autocast12import cv213from matplotlib import pyplot as plt14from torchvision import transforms15from diffusers import DiffusionPipeline16from diffusers.utils import torch_device17pipe = DiffusionPipeline.from_pretrained(18 "Fantasy-Studio/Paint-by-Example",19 torch_dtype=torch.float16,20)21pipe = pipe.to("cuda")22 23from share_btn import community_icon_html, loading_icon_html, share_js24 25def read_content(file_path: str) -> str:26 """read the content of target file27 """28 with open(file_path, 'r', encoding='utf-8') as f:29 content = f.read()30 31 return content32 33def predict(dict, reference, scale, seed, step):34 width,height=dict["image"].size35 if width<height:36 factor=width/512.037 width=51238 height=int((height/factor)/8.0)*839 40 else:41 factor=height/512.042 height=51243 width=int((width/factor)/8.0)*844 init_image = dict["image"].convert("RGB").resize((width,height))45 mask = dict["mask"].convert("RGB").resize((width,height))46 generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None47 output = pipe(48 image=init_image,49 mask_image=mask,50 example_image=reference,51 generator=generator,52 guidance_scale=scale,53 num_inference_steps=step,54 ).images[0]55 return output, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)56 57 58css = '''59.container {max-width: 1150px;margin: auto;padding-top: 1.5rem}60#image_upload{min-height:400px}61#image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 400px}62#mask_radio .gr-form{background:transparent; border: none}63#word_mask{margin-top: .75em !important}64#word_mask textarea:disabled{opacity: 0.3}65.footer {margin-bottom: 45px;margin-top: 35px;text-align: center;border-bottom: 1px solid #e5e5e5}66.footer>p {font-size: .8rem; display: inline-block; padding: 0 10px;transform: translateY(10px);background: white}67.dark .footer {border-color: #303030}68.dark .footer>p {background: #0b0f19}69.acknowledgments h4{margin: 1.25em 0 .25em 0;font-weight: bold;font-size: 115%}70#image_upload .touch-none{display: flex}71@keyframes spin {72 from {73 transform: rotate(0deg);74 }75 to {76 transform: rotate(360deg);77 }78}79#share-btn-container {80 display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;81}82#share-btn {83 all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;84}85#share-btn * {86 all: unset;87}88#share-btn-container div:nth-child(-n+2){89 width: auto !important;90 min-height: 0px !important;91}92#share-btn-container .wrap {93 display: none !important;94}95'''96example={}97ref_dir='examples/reference'98image_dir='examples/image'99ref_list=[os.path.join(ref_dir,file) for file in os.listdir(ref_dir)]100ref_list.sort()101image_list=[os.path.join(image_dir,file) for file in os.listdir(image_dir)]102image_list.sort()103 104 105image_blocks = gr.Blocks(css=css)106with image_blocks as demo:107 gr.HTML(read_content("header.html"))108 with gr.Group():109 with gr.Box():110 with gr.Row():111 with gr.Column():112 image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Source Image")113 reference = gr.Image(source='upload', elem_id="image_upload", type="pil", label="Reference Image")114 115 with gr.Column():116 image_out = gr.Image(label="Output", elem_id="output-img").style(height=400)117 guidance = gr.Slider(label="Guidance scale", value=5, maximum=15,interactive=True)118 steps = gr.Slider(label="Steps", value=50, minimum=2, maximum=75, step=1,interactive=True)119 120 seed = gr.Slider(0, 10000, label='Seed (0 = random)', value=0, step=1)121 122 with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):123 btn = gr.Button("Paint!").style(124 margin=False,125 rounded=(False, True, True, False),126 full_width=True,127 )128 with gr.Group(elem_id="share-btn-container"):129 community_icon = gr.HTML(community_icon_html, visible=True)130 loading_icon = gr.HTML(loading_icon_html, visible=True)131 share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)132 133 134 with gr.Row():135 with gr.Column():136 gr.Examples(image_list, inputs=[image],label="Examples - Source Image",examples_per_page=12)137 with gr.Column():138 gr.Examples(ref_list, inputs=[reference],label="Examples - Reference Image",examples_per_page=12)139 140 btn.click(fn=predict, inputs=[image, reference, guidance, seed, steps], outputs=[image_out, community_icon, loading_icon, share_button])141 share_button.click(None, [], [], _js=share_js)142 143 144 145 gr.HTML(146 """147 <div class="footer">148 <p>Model by <a href="" style="text-decoration: underline;" target="_blank">Fantasy-Studio</a> - Gradio Demo by 🤗 Hugging Face149 </p>150 </div>151 <div class="acknowledgments">152 <p><h4>LICENSE</h4>153 The model is licensed with a <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" style="text-decoration: underline;" target="_blank">CreativeML Open RAIL-M</a> license. The authors claim no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in this license. The license forbids you from sharing any content that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. For the full list of restrictions please <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" target="_blank" style="text-decoration: underline;" target="_blank">read the license</a></p>154 """155 )156 157image_blocks.launch()