Toar-HF/diffusers-image
0
1import gradio as gr2import spaces3import torch4from diffusers import AutoencoderKL, TCDScheduler5from diffusers.models.model_loading_utils import load_state_dict6from gradio_imageslider import ImageSlider7from huggingface_hub import hf_hub_download8 9from controlnet_union import ControlNetModel_Union10from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline11 12from PIL import Image, ImageDraw13import numpy as np14 15config_file = hf_hub_download(16 "xinsir/controlnet-union-sdxl-1.0",17 filename="config_promax.json",18)19 20config = ControlNetModel_Union.load_config(config_file)21controlnet_model = ControlNetModel_Union.from_config(config)22model_file = hf_hub_download(23 "xinsir/controlnet-union-sdxl-1.0",24 filename="diffusion_pytorch_model_promax.safetensors",25)26state_dict = load_state_dict(model_file)27model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(28 controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"29)30model.to(device="cuda", dtype=torch.float16)31 32vae = AutoencoderKL.from_pretrained(33 "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float1634).to("cuda")35 36pipe = StableDiffusionXLFillPipeline.from_pretrained(37 "SG161222/RealVisXL_V5.0_Lightning",38 torch_dtype=torch.float16,39 vae=vae,40 controlnet=model,41 variant="fp16",42).to("cuda")43 44pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)45 46 47def can_expand(source_width, source_height, target_width, target_height, alignment):48 """Checks if the image can be expanded based on the alignment."""49 if alignment in ("Left", "Right") and source_width >= target_width:50 return False51 if alignment in ("Top", "Bottom") and source_height >= target_height:52 return False53 return True54 55def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):56 target_size = (width, height)57 58 # Calculate the scaling factor to fit the image within the target size59 scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)60 new_width = int(image.width * scale_factor)61 new_height = int(image.height * scale_factor)62 63 # Resize the source image to fit within target size64 source = image.resize((new_width, new_height), Image.LANCZOS)65 66 # Apply resize option using percentages67 if resize_option == "Full":68 resize_percentage = 10069 elif resize_option == "50%":70 resize_percentage = 5071 elif resize_option == "33%":72 resize_percentage = 3373 elif resize_option == "25%":74 resize_percentage = 2575 else: # Custom76 resize_percentage = custom_resize_percentage77 78 # Calculate new dimensions based on percentage79 resize_factor = resize_percentage / 10080 new_width = int(source.width * resize_factor)81 new_height = int(source.height * resize_factor)82 83 # Ensure minimum size of 64 pixels84 new_width = max(new_width, 64)85 new_height = max(new_height, 64)86 87 # Resize the image88 source = source.resize((new_width, new_height), Image.LANCZOS)89 90 # Calculate the overlap in pixels based on the percentage91 overlap_x = int(new_width * (overlap_percentage / 100))92 overlap_y = int(new_height * (overlap_percentage / 100))93 94 # Ensure minimum overlap of 1 pixel95 overlap_x = max(overlap_x, 1)96 overlap_y = max(overlap_y, 1)97 98 # Calculate margins based on alignment99 if alignment == "Middle":100 margin_x = (target_size[0] - new_width) // 2101 margin_y = (target_size[1] - new_height) // 2102 elif alignment == "Left":103 margin_x = 0104 margin_y = (target_size[1] - new_height) // 2105 elif alignment == "Right":106 margin_x = target_size[0] - new_width107 margin_y = (target_size[1] - new_height) // 2108 elif alignment == "Top":109 margin_x = (target_size[0] - new_width) // 2110 margin_y = 0111 elif alignment == "Bottom":112 margin_x = (target_size[0] - new_width) // 2113 margin_y = target_size[1] - new_height114 115 # Adjust margins to eliminate gaps116 margin_x = max(0, min(margin_x, target_size[0] - new_width))117 margin_y = max(0, min(margin_y, target_size[1] - new_height))118 119 # Create a new background image and paste the resized source image120 background = Image.new('RGB', target_size, (255, 255, 255))121 background.paste(source, (margin_x, margin_y))122 123 # Create the mask124 mask = Image.new('L', target_size, 255)125 mask_draw = ImageDraw.Draw(mask)126 127 # Calculate overlap areas128 white_gaps_patch = 2129 130 left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch131 right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch132 top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch133 bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch134 135 if alignment == "Left":136 left_overlap = margin_x + overlap_x if overlap_left else margin_x137 elif alignment == "Right":138 right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width139 elif alignment == "Top":140 top_overlap = margin_y + overlap_y if overlap_top else margin_y141 elif alignment == "Bottom":142 bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height143 144 145 # Draw the mask146 mask_draw.rectangle([147 (left_overlap, top_overlap),148 (right_overlap, bottom_overlap)149 ], fill=0)150 151 return background, mask152 153def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):154 background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)155 156 # Create a preview image showing the mask157 preview = background.copy().convert('RGBA')158 159 # Create a semi-transparent red overlay160 red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64)) # Reduced alpha to 64 (25% opacity)161 162 # Convert black pixels in the mask to semi-transparent red163 red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))164 red_mask.paste(red_overlay, (0, 0), mask)165 166 # Overlay the red mask on the background167 preview = Image.alpha_composite(preview, red_mask)168 169 return preview170 171@spaces.GPU(duration=24)172def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):173 background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)174 175 if not can_expand(background.width, background.height, width, height, alignment):176 alignment = "Middle"177 178 cnet_image = background.copy()179 cnet_image.paste(0, (0, 0), mask)180 181 final_prompt = f"{prompt_input} , high quality, 4k"182 183 (184 prompt_embeds,185 negative_prompt_embeds,186 pooled_prompt_embeds,187 negative_pooled_prompt_embeds,188 ) = pipe.encode_prompt(final_prompt, "cuda", True)189 190 for image in pipe(191 prompt_embeds=prompt_embeds,192 negative_prompt_embeds=negative_prompt_embeds,193 pooled_prompt_embeds=pooled_prompt_embeds,194 negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,195 image=cnet_image,196 num_inference_steps=num_inference_steps197 ):198 yield cnet_image, image199 200 image = image.convert("RGBA")201 cnet_image.paste(image, (0, 0), mask)202 203 yield background, cnet_image204 205def clear_result():206 """Clears the result ImageSlider."""207 return gr.update(value=None)208 209def preload_presets(target_ratio, ui_width, ui_height):210 """Updates the width and height sliders based on the selected aspect ratio."""211 if target_ratio == "9:16":212 changed_width = 720213 changed_height = 1280214 return changed_width, changed_height, gr.update()215 elif target_ratio == "16:9":216 changed_width = 1280217 changed_height = 720218 return changed_width, changed_height, gr.update()219 elif target_ratio == "1:1":220 changed_width = 1024221 changed_height = 1024222 return changed_width, changed_height, gr.update()223 elif target_ratio == "Custom":224 return ui_width, ui_height, gr.update(open=True)225 226def select_the_right_preset(user_width, user_height):227 if user_width == 720 and user_height == 1280:228 return "9:16"229 elif user_width == 1280 and user_height == 720:230 return "16:9"231 elif user_width == 1024 and user_height == 1024:232 return "1:1"233 else:234 return "Custom"235 236def toggle_custom_resize_slider(resize_option):237 return gr.update(visible=(resize_option == "Custom"))238 239def update_history(new_image, history):240 """Updates the history gallery with the new image."""241 if history is None:242 history = []243 history.insert(0, new_image)244 return history245 246css = """247.gradio-container {248 width: 1200px !important;249}250"""251 252title = """<h1 align="center">Diffusers Image Outpaint</h1>253<div align="center">Drop an image you would like to extend, pick your expected ratio and hit Generate.</div>254<div style="display: flex; justify-content: center; align-items: center; text-align: center;">255 <p style="display: flex;gap: 6px;">256 <a href="https://huggingface.co/spaces/fffiloni/diffusers-image-outpaint?duplicate=true">257 <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate this Space">258 </a> to skip the queue and enjoy faster inference on the GPU of your choice 259 </p>260</div>261"""262 263with gr.Blocks(css=css) as demo:264 with gr.Column():265 gr.HTML(title)266 267 with gr.Row():268 with gr.Column():269 input_image = gr.Image(270 type="pil",271 label="Input Image"272 )273 274 with gr.Row():275 with gr.Column(scale=2):276 prompt_input = gr.Textbox(label="Prompt (Optional)")277 with gr.Column(scale=1):278 run_button = gr.Button("Generate")279 280 with gr.Row():281 target_ratio = gr.Radio(282 label="Expected Ratio",283 choices=["9:16", "16:9", "1:1", "Custom"],284 value="9:16",285 scale=2286 )287 288 alignment_dropdown = gr.Dropdown(289 choices=["Middle", "Left", "Right", "Top", "Bottom"],290 value="Middle",291 label="Alignment"292 )293 294 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:295 with gr.Column():296 with gr.Row():297 width_slider = gr.Slider(298 label="Target Width",299 minimum=720,300 maximum=1536,301 step=8,302 value=720, # Set a default value303 )304 height_slider = gr.Slider(305 label="Target Height",306 minimum=720,307 maximum=1536,308 step=8,309 value=1280, # Set a default value310 )311 312 num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)313 with gr.Group():314 overlap_percentage = gr.Slider(315 label="Mask overlap (%)",316 minimum=1,317 maximum=50,318 value=10,319 step=1320 )321 with gr.Row():322 overlap_top = gr.Checkbox(label="Overlap Top", value=True)323 overlap_right = gr.Checkbox(label="Overlap Right", value=True)324 with gr.Row():325 overlap_left = gr.Checkbox(label="Overlap Left", value=True)326 overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)327 with gr.Row():328 resize_option = gr.Radio(329 label="Resize input image",330 choices=["Full", "50%", "33%", "25%", "Custom"],331 value="Full"332 )333 custom_resize_percentage = gr.Slider(334 label="Custom resize (%)",335 minimum=1,336 maximum=100,337 step=1,338 value=50,339 visible=False340 )341 342 with gr.Column():343 preview_button = gr.Button("Preview alignment and mask")344 345 346 gr.Examples(347 examples=[348 ["./examples/example_1.webp", 1280, 720, "Middle"],349 ["./examples/example_2.jpg", 1440, 810, "Left"],350 ["./examples/example_3.jpg", 1024, 1024, "Top"],351 ["./examples/example_3.jpg", 1024, 1024, "Bottom"],352 ],353 inputs=[input_image, width_slider, height_slider, alignment_dropdown],354 )355 356 357 358 with gr.Column():359 result = ImageSlider(360 interactive=False,361 label="Generated Image",362 )363 use_as_input_button = gr.Button("Use as Input Image", visible=False)364 365 history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)366 preview_image = gr.Image(label="Preview")367 368 369 370 def use_output_as_input(output_image):371 """Sets the generated output as the new input image."""372 return gr.update(value=output_image[1])373 374 use_as_input_button.click(375 fn=use_output_as_input,376 inputs=[result],377 outputs=[input_image]378 )379 380 target_ratio.change(381 fn=preload_presets,382 inputs=[target_ratio, width_slider, height_slider],383 outputs=[width_slider, height_slider, settings_panel],384 queue=False385 )386 387 width_slider.change(388 fn=select_the_right_preset,389 inputs=[width_slider, height_slider],390 outputs=[target_ratio],391 queue=False392 )393 394 height_slider.change(395 fn=select_the_right_preset,396 inputs=[width_slider, height_slider],397 outputs=[target_ratio],398 queue=False399 )400 401 resize_option.change(402 fn=toggle_custom_resize_slider,403 inputs=[resize_option],404 outputs=[custom_resize_percentage],405 queue=False406 )407 408 run_button.click( # Clear the result409 fn=clear_result,410 inputs=None,411 outputs=result,412 ).then( # Generate the new image413 fn=infer,414 inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,415 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,416 overlap_left, overlap_right, overlap_top, overlap_bottom],417 outputs=result,418 ).then( # Update the history gallery419 fn=lambda x, history: update_history(x[1], history),420 inputs=[result, history_gallery],421 outputs=history_gallery,422 ).then( # Show the "Use as Input Image" button423 fn=lambda: gr.update(visible=True),424 inputs=None,425 outputs=use_as_input_button,426 )427 428 prompt_input.submit( # Clear the result429 fn=clear_result,430 inputs=None,431 outputs=result,432 ).then( # Generate the new image433 fn=infer,434 inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,435 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,436 overlap_left, overlap_right, overlap_top, overlap_bottom],437 outputs=result,438 ).then( # Update the history gallery439 fn=lambda x, history: update_history(x[1], history),440 inputs=[result, history_gallery],441 outputs=history_gallery,442 ).then( # Show the "Use as Input Image" button443 fn=lambda: gr.update(visible=True),444 inputs=None,445 outputs=use_as_input_button,446 )447 448 preview_button.click(449 fn=preview_image_and_mask,450 inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option, custom_resize_percentage, alignment_dropdown,451 overlap_left, overlap_right, overlap_top, overlap_bottom],452 outputs=preview_image,453 queue=False454 )455 456demo.queue(max_size=12).launch(share=False, show_error=True)