import gradio as gr import spaces import torch from diffusers import AutoencoderKL, TCDScheduler # (Assume ControlNet manual load or from_pretrained is already working) from controlnet_union import ControlNetModel_Union from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline from gradio_imageslider import ImageSlider from huggingface_hub import hf_hub_download from PIL import Image, ImageDraw import numpy as np # --- Load ControlNet and SDXL Fill Pipeline --- # (Either manual download or via from_pretrained) controlnet_model = ControlNetModel_Union.from_pretrained( "xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16, variant="fp16" ).to("cuda") vae = AutoencoderKL.from_pretrained( "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 ).to("cuda") pipe = StableDiffusionXLFillPipeline.from_pretrained( "SG161222/RealVisXL_V5.0_Lightning", torch_dtype=torch.float16, vae=vae, controlnet=controlnet_model, variant="fp16", ).to("cuda") pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config) # --- Utility functions --- def can_expand(source_width, source_height, target_width, target_height, alignment): if alignment in ("Left", "Right") and source_width >= target_width: return False if alignment in ("Top", "Bottom") and source_height >= target_height: return False return True def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom): target = (width, height) scale = min(target[0] / image.width, target[1] / image.height) w, h = int(image.width * scale), int(image.height * scale) src = image.resize((w, h), Image.LANCZOS) # Resize percentage if resize_option == "Full": pct = 100 elif resize_option == "50%": pct = 50 elif resize_option == "33%": pct = 33 elif resize_option == "25%": pct = 25 else: pct = custom_resize_percentage rw, rh = max(int(src.width * pct / 100), 64), max(int(src.height * pct / 100), 64) src = src.resize((rw, rh), Image.LANCZOS) ox = max(int(rw * overlap_percentage / 100), 1) oy = max(int(rh * overlap_percentage / 100), 1) # Margins if alignment == "Middle": mx, my = (width - rw)//2, (height - rh)//2 elif alignment == "Left": mx, my = 0, (height - rh)//2 elif alignment == "Right": mx, my = width - rw, (height - rh)//2 elif alignment == "Top": mx, my = (width - rw)//2, 0 else: mx, my = (width - rw)//2, height - rh mx, my = max(0, min(mx, width - rw)), max(0, min(my, height - rh)) bg = Image.new("RGB", target, (255,255,255)) bg.paste(src, (mx, my)) mask = Image.new("L", target, 255) d = ImageDraw.Draw(mask) lx = mx + (ox if overlap_left else 2) rx = mx + rw - (ox if overlap_right else 2) ty = my + (oy if overlap_top else 2) by = my + rh - (oy if overlap_bottom else 2) # Edge adjustments if alignment == "Left": lx = mx + (ox if overlap_left else 0) if alignment == "Right": rx = mx + rw - (ox if overlap_right else 0) if alignment == "Top": ty = my + (oy if overlap_top else 0) if alignment == "Bottom": by = my + rh - (oy if overlap_bottom else 0) d.rectangle([(lx, ty), (rx, by)], fill=0) return bg, mask def preview_image_and_mask(*args): bg, mask = prepare_image_and_mask(*args) vis = bg.copy().convert("RGBA") red = Image.new("RGBA", bg.size, (255,0,0,64)) overlay = Image.new("RGBA", bg.size, (0,0,0,0)) overlay.paste(red, (0,0), mask) return Image.alpha_composite(vis, overlay) # --- Fixed infer: return list for slider --- @spaces.GPU(duration=24) def 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): 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 ) if not can_expand(background.width, background.height, width, height, alignment): alignment = "Middle" hole = background.copy() hole.paste(0, (0,0), mask) final_prompt = f"{prompt_input} , high quality, 4k" embeds = pipe.encode_prompt(final_prompt, "cuda", True) # Run pipeline and grab last frame gen = pipe( prompt_embeds=embeds[0], negative_prompt_embeds=embeds[1], pooled_prompt_embeds=embeds[2], negative_pooled_prompt_embeds=embeds[3], image=hole, num_inference_steps=num_inference_steps ) last = None for img in gen: last = img out = last.convert("RGBA") hole.paste(out, (0,0), mask) # Return a list: [input_hole_image, final_output] return [background, hole] def clear_result(): return gr.update(value=None) def preload_presets(ratio, w, h): if ratio == "9:16": return 720, 1280, gr.update() if ratio == "16:9": return 1280, 720, gr.update() if ratio == "1:1": return 1024, 1024, gr.update() return w, h, gr.update(open=True) def select_the_right_preset(w, h): if (w,h) == (720,1280): return "9:16" if (w,h) == (1280,720): return "16:9" if (w,h) == (1024,1024): return "1:1" return "Custom" def toggle_custom_resize_slider(opt): return gr.update(visible=(opt=="Custom")) def update_history(img, history): history = history or [] history.insert(0, img) return history css = ".gradio-container { width: 1200px !important; }" title = "

Diffusers Image Outpaint Lightning

" with gr.Blocks(css=css) as demo: gr.HTML(title) with gr.Row(): with gr.Column(): input_image = gr.Image(type="pil", label="Input Image") prompt_input = gr.Textbox(label="Prompt (Optional)") run_button = gr.Button("Generate") target_ratio = gr.Radio(["9:16","16:9","1:1","Custom"], value="9:16", label="Expected Ratio") alignment_dropdown = gr.Dropdown(["Middle","Left","Right","Top","Bottom"], value="Middle", label="Alignment") with gr.Accordion("Advanced settings", open=False) as adv: width_slider = gr.Slider(720,1536,step=8, value=720, label="Target Width") height_slider = gr.Slider(720,1536,step=8, value=1280, label="Target Height") num_steps = gr.Slider(4,12,step=1, value=8, label="Steps") overlap_pct = gr.Slider(1,50,step=1, value=10, label="Mask overlap (%)") overlap_top = gr.Checkbox(label="Overlap Top", value=True) overlap_right = gr.Checkbox(label="Overlap Right", value=True) overlap_left = gr.Checkbox(label="Overlap Left", value=True) overlap_bottom= gr.Checkbox(label="Overlap Bottom", value=True) resize_opt = gr.Radio(["Full","50%","33%","25%","Custom"], value="Full", label="Resize input image") custom_resize = gr.Slider(1,100,step=1, value=50, visible=False, label="Custom resize (%)") preview_btn = gr.Button("Preview alignment and mask") gr.Examples( examples=[ ["./examples/example_1.webp",1280,720,"Middle"], ["./examples/example_2.jpg",1440,810,"Left"], ["./examples/example_3.jpg",1024,1024,"Top"], ["./examples/example_3.jpg",1024,1024,"Bottom"] ], inputs=[input_image,width_slider,height_slider,alignment_dropdown] ) with gr.Column(): result = ImageSlider(label="Comparison", interactive=False, type="pil", slider_color="pink") history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain") preview_image = gr.Image(label="Preview") # Callbacks run_button.click(clear_result, None, result) run_button.click( infer, inputs=[ input_image, width_slider, height_slider, overlap_pct, num_steps, resize_opt, custom_resize, prompt_input, alignment_dropdown, overlap_left, overlap_right, overlap_top, overlap_bottom], outputs=result ).then(update_history, inputs=[result, history_gallery], outputs=history_gallery) target_ratio.change(preload_presets, [target_ratio, width_slider, height_slider], [width_slider, height_slider, adv]) width_slider.change(select_the_right_preset, [width_slider, height_slider], target_ratio) height_slider.change(select_the_right_preset, [width_slider, height_slider], target_ratio) resize_opt.change(toggle_custom_resize_slider, resize_opt, custom_resize) preview_btn.click(preview_image_and_mask, [input_image, width_slider, height_slider, overlap_pct, resize_opt, custom_resize, alignment_dropdown, overlap_left, overlap_right, overlap_top, overlap_bottom], preview_image) demo.queue(max_size=20).launch(share=False, ssr_mode=False, show_error=True)