# -*- coding: utf-8 -*-
import gradio as gr
import spaces
import torch
from diffusers import AutoencoderKL, TCDScheduler
from diffusers.models.model_loading_utils import load_state_dict
# Remove ImageSlider import as it's no longer needed
# from gradio_imageslider import ImageSlider
from huggingface_hub import hf_hub_download
from controlnet_union import ControlNetModel_Union
from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
from PIL import Image, ImageDraw
import numpy as np
# --- Model Loading (Keep as is) ---
config_file = hf_hub_download(
"xinsir/controlnet-union-sdxl-1.0",
filename="config_promax.json",
)
config = ControlNetModel_Union.load_config(config_file)
controlnet_model = ControlNetModel_Union.from_config(config)
model_file = hf_hub_download(
"xinsir/controlnet-union-sdxl-1.0",
filename="diffusion_pytorch_model_promax.safetensors",
)
state_dict = load_state_dict(model_file)
model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
)
model.to(device="cuda", dtype=torch.float16)
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=model,
variant="fp16",
).to("cuda")
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
# --- Helper Functions (Keep as is, except infer) ---
def can_expand(source_width, source_height, target_width, target_height, alignment):
"""Checks if the image can be expanded based on the 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_size = (width, height)
# Calculate the scaling factor to fit the image within the target size
scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
new_width = int(image.width * scale_factor)
new_height = int(image.height * scale_factor)
# Resize the source image to fit within target size
source = image.resize((new_width, new_height), Image.LANCZOS)
# Apply resize option using percentages
if resize_option == "Full":
resize_percentage = 100
elif resize_option == "50%":
resize_percentage = 50
elif resize_option == "33%":
resize_percentage = 33
elif resize_option == "25%":
resize_percentage = 25
else: # Custom
resize_percentage = custom_resize_percentage
# Calculate new dimensions based on percentage
resize_factor = resize_percentage / 100
new_width = int(source.width * resize_factor)
new_height = int(source.height * resize_factor)
# Ensure minimum size of 64 pixels
new_width = max(new_width, 64)
new_height = max(new_height, 64)
# Resize the image
source = source.resize((new_width, new_height), Image.LANCZOS)
# Calculate the overlap in pixels based on the percentage
overlap_x = int(new_width * (overlap_percentage / 100))
overlap_y = int(new_height * (overlap_percentage / 100))
# Ensure minimum overlap of 1 pixel
overlap_x = max(overlap_x, 1)
overlap_y = max(overlap_y, 1)
# Calculate margins based on alignment
if alignment == "Middle":
margin_x = (target_size[0] - new_width) // 2
margin_y = (target_size[1] - new_height) // 2
elif alignment == "Left":
margin_x = 0
margin_y = (target_size[1] - new_height) // 2
elif alignment == "Right":
margin_x = target_size[0] - new_width
margin_y = (target_size[1] - new_height) // 2
elif alignment == "Top":
margin_x = (target_size[0] - new_width) // 2
margin_y = 0
elif alignment == "Bottom":
margin_x = (target_size[0] - new_width) // 2
margin_y = target_size[1] - new_height
# Adjust margins to eliminate gaps
margin_x = max(0, min(margin_x, target_size[0] - new_width))
margin_y = max(0, min(margin_y, target_size[1] - new_height))
# Create a new background image and paste the resized source image
background = Image.new('RGB', target_size, (255, 255, 255))
background.paste(source, (margin_x, margin_y))
# Create the mask
mask = Image.new('L', target_size, 255)
mask_draw = ImageDraw.Draw(mask)
# Calculate overlap areas
white_gaps_patch = 2
left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
if alignment == "Left":
left_overlap = margin_x + overlap_x if overlap_left else margin_x
elif alignment == "Right":
right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
elif alignment == "Top":
top_overlap = margin_y + overlap_y if overlap_top else margin_y
elif alignment == "Bottom":
bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
# Draw the mask
mask_draw.rectangle([
(left_overlap, top_overlap),
(right_overlap, bottom_overlap)
], fill=0)
return background, mask
def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, 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)
# Create a preview image showing the mask
preview = background.copy().convert('RGBA')
# Create a semi-transparent red overlay
red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64)) # Reduced alpha to 64 (25% opacity)
# Convert black pixels in the mask to semi-transparent red
red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
red_mask.paste(red_overlay, (0, 0), mask)
# Overlay the red mask on the background
preview = Image.alpha_composite(preview, red_mask)
return preview
@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):
if image is None:
raise gr.Error("Please upload an input image.")
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):
# Optionally provide feedback or default to middle
# gr.Warning(f"Cannot expand image with '{alignment}' alignment as source dimension is larger than target. Defaulting to 'Middle'.")
alignment = "Middle"
# Recalculate background and mask if alignment changed due to this check
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)
cnet_image = background.copy()
# Apply mask to create the input for controlnet (black out non-masked area)
# cnet_image.paste(0, (0, 0), mask) # This line seems incorrect for inpainting/outpainting, usually the unmasked area is kept
# The pipeline expects the original image content where mask=0 and potentially noise/latents where mask=1
# Let's keep the original image content in the unmasked area and let the pipeline handle the masked area.
# The `StableDiffusionXLFillPipeline` likely uses the `image` input and `mask` differently than standard inpainting.
# Based on typical diffusers pipelines, `image` is often the *original* content placed on the canvas.
# Let's pass `background` as the image input for the pipeline.
final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = pipe.encode_prompt(final_prompt, "cuda", True, negative_prompt="") # Add default negative prompt
# The pipeline expects the `image` and `mask_image` arguments for inpainting/outpainting
# `image` should be the canvas with the original image placed.
# `mask_image` defines the area to be filled (white=fill, black=keep).
# Our mask is inverted (black=keep, white=fill). Invert it.
inverted_mask = Image.fromarray(255 - np.array(mask))
# Run the pipeline
# Note: The generator inside the pipeline call is not used here as we only need the final result.
# We iterate once to get the final image.
generated_image = None
for img_output in pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
image=background, # Pass the background with the source image placed
mask_image=inverted_mask, # Pass the inverted mask (white = area to fill)
control_image=background, # ControlNet Union might need the full image context
num_inference_steps=num_inference_steps,
output_type="pil" # Ensure PIL images are returned
):
generated_image = img_output[0] # The pipeline returns a list containing the image
if generated_image is None:
raise gr.Error("Image generation failed.")
# The pipeline should return the complete image already composited.
# No need to manually paste.
final_image = generated_image.convert("RGB")
# Yield only the final generated image
yield final_image
def clear_result():
"""Clears the result Image component."""
return gr.update(value=None)
def preload_presets(target_ratio, ui_width, ui_height):
"""Updates the width and height sliders based on the selected aspect ratio."""
if target_ratio == "9:16":
changed_width = 720
changed_height = 1280
return changed_width, changed_height, gr.update(open=False) # Close accordion
elif target_ratio == "16:9":
changed_width = 1280
changed_height = 720
return changed_width, changed_height, gr.update(open=False) # Close accordion
elif target_ratio == "1:1":
changed_width = 1024
changed_height = 1024
return changed_width, changed_height, gr.update(open=False) # Close accordion
elif target_ratio == "Custom":
# Keep current slider values but open the accordion
return ui_width, ui_height, gr.update(open=True)
def select_the_right_preset(user_width, user_height):
"""Selects the preset radio button based on current width/height."""
if user_width == 720 and user_height == 1280:
return "9:16"
elif user_width == 1280 and user_height == 720:
return "16:9"
elif user_width == 1024 and user_height == 1024:
return "1:1"
else:
return "Custom"
def toggle_custom_resize_slider(resize_option):
"""Shows/hides the custom resize slider."""
return gr.update(visible=(resize_option == "Custom"))
def update_history(new_image, history):
"""Updates the history gallery with the new image."""
if new_image is None: # Don't add None to history
return history
if history is None:
history = []
# Prepend the new image (as PIL) to the history list
history.insert(0, new_image)
# Limit history size if desired (e.g., keep last 12)
max_history = 12
if len(history) > max_history:
history = history[:max_history]
return history
# --- Gradio UI ---
css = """
.gradio-container {
max-width: 1200px !important; /* Limit overall width */
margin: auto; /* Center the container */
}
/* Ensure gallery items are reasonably sized */
#history_gallery .thumbnail-item {
height: 100px !important; /* Adjust as needed */
}
#history_gallery .gallery {
grid-template-columns: repeat(auto-fill, minmax(100px, 1fr)) !important; /* Adjust column size */
}
"""
title = """Diffusers Image Outpaint
Drop an image you would like to extend, pick your expected ratio and hit Generate.
to skip the queue and enjoy faster inference on the GPU of your choice
"""
with gr.Blocks(css=css) as demo:
with gr.Column():
gr.HTML(title)
with gr.Row():
with gr.Column(scale=1): # Input column
input_image = gr.Image(
type="pil",
label="Input Image"
)
with gr.Row():
with gr.Column(scale=2):
prompt_input = gr.Textbox(label="Prompt (Optional)", placeholder="Describe the desired extended scene...")
with gr.Column(scale=1, min_width=150):
run_button = gr.Button("Generate", variant="primary")
with gr.Row():
target_ratio = gr.Radio(
label="Target Ratio",
choices=["9:16", "16:9", "1:1", "Custom"],
value="9:16",
scale=2
)
alignment_dropdown = gr.Dropdown(
choices=["Middle", "Left", "Right", "Top", "Bottom"],
value="Middle",
label="Align Source Image"
)
with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
with gr.Column():
with gr.Row():
width_slider = gr.Slider(
label="Target Width (px)",
minimum=512, # Lowered min slightly
maximum=2048, # Increased max slightly
step=64, # SDXL optimal step
value=720,
)
height_slider = gr.Slider(
label="Target Height (px)",
minimum=512, # Lowered min slightly
maximum=2048, # Increased max slightly
step=64, # SDXL optimal step
value=1280,
)
num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=20, step=1, value=8) # Increased max steps slightly
with gr.Group():
overlap_percentage = gr.Slider(
label="Mask overlap (%)",
minimum=1,
maximum=50,
value=10,
step=1,
info="How much the new area overlaps the original image."
)
gr.Markdown("Select sides to overlap (influences mask generation):")
with gr.Row():
overlap_top = gr.Checkbox(label="Top", value=True)
overlap_right = gr.Checkbox(label="Right", value=True)
with gr.Row():
overlap_left = gr.Checkbox(label="Left", value=True)
overlap_bottom = gr.Checkbox(label="Bottom", value=True)
with gr.Row():
resize_option = gr.Radio(
label="Resize input image before placing",
choices=["Full", "50%", "33%", "25%", "Custom"],
value="Full",
info="Scales the source image down before placing it on the target canvas."
)
custom_resize_percentage = gr.Slider(
label="Custom resize (%)",
minimum=1,
maximum=100,
step=1,
value=50,
visible=False
)
with gr.Column():
preview_button = gr.Button("Preview Alignment & Mask")
gr.Examples(
examples=[
["./examples/example_1.webp", 1280, 720, "Middle", "A wide landscape view of the mountains"],
["./examples/example_2.jpg", 1440, 810, "Left", "Full body shot of the cat sitting on a rug"],
["./examples/example_3.jpg", 1024, 1024, "Top", "The cloudy sky above the building"],
["./examples/example_3.jpg", 1024, 1024, "Bottom", "The street below the building"],
],
inputs=[input_image, width_slider, height_slider, alignment_dropdown, prompt_input],
label="Examples (Click to load)"
)
with gr.Column(scale=1): # Output column
# Replace ImageSlider with gr.Image
result_image = gr.Image(
label="Generated Image",
interactive=False,
show_download_button=True,
type="pil" # Ensure output is PIL for history
)
with gr.Row():
use_as_input_button = gr.Button("Use as Input", visible=False)
clear_button = gr.Button("Clear Output") # Added clear button
preview_mask_image = gr.Image(label="Alignment & Mask Preview", interactive=False)
history_gallery = gr.Gallery(
label="History",
columns=4, # Adjust columns as needed
object_fit="contain",
interactive=False,
show_label=True,
elem_id="history_gallery",
height=300 # Set a fixed height for the gallery area
)
# --- Event Handlers ---
def use_output_as_input(output_pil_image):
"""Sets the generated output PIL image as the new input image."""
# output_image comes directly from result_image which is PIL type
return gr.update(value=output_pil_image)
use_as_input_button.click(
fn=use_output_as_input,
inputs=[result_image], # Input is the single result image
outputs=[input_image]
)
clear_button.click(
fn=lambda: (gr.update(value=None), gr.update(visible=False), gr.update(value=None)), # Clear image, hide button, clear preview
inputs=None,
outputs=[result_image, use_as_input_button, preview_mask_image],
queue=False
)
target_ratio.change(
fn=preload_presets,
inputs=[target_ratio, width_slider, height_slider],
outputs=[width_slider, height_slider, settings_panel],
queue=False
)
# Link sliders back to ratio selector and potentially open accordion
width_slider.change(
fn=lambda w, h: (select_the_right_preset(w, h), gr.update() if select_the_right_preset(w, h) == "Custom" else gr.update()),
inputs=[width_slider, height_slider],
outputs=[target_ratio, settings_panel],
queue=False
)
height_slider.change(
fn=lambda w, h: (select_the_right_preset(w, h), gr.update() if select_the_right_preset(w, h) == "Custom" else gr.update()),
inputs=[width_slider, height_slider],
outputs=[target_ratio, settings_panel],
queue=False
)
resize_option.change(
fn=toggle_custom_resize_slider,
inputs=[resize_option],
outputs=[custom_resize_percentage],
queue=False
)
# Define common inputs for generation
gen_inputs = [
input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
overlap_left, overlap_right, overlap_top, overlap_bottom
]
# Define common steps after generation
def handle_output(generated_image, current_history):
# generated_image is the single PIL image from infer
new_history = update_history(generated_image, current_history)
button_visibility = gr.update(visible=True) if generated_image else gr.update(visible=False)
return generated_image, new_history, button_visibility
run_button.click(
fn=lambda: (gr.update(value=None), gr.update(visible=False)), # Clear result and hide button first
inputs=None,
outputs=[result_image, use_as_input_button],
queue=False # Don't queue the clearing part
).then(
fn=infer, # Run the generation
inputs=gen_inputs,
outputs=result_image, # Output is the single generated image
).then(
fn=handle_output, # Process output: update history, show button
inputs=[result_image, history_gallery],
outputs=[result_image, history_gallery, use_as_input_button] # Update result again (no change), history, and button visibility
)
prompt_input.submit(
fn=lambda: (gr.update(value=None), gr.update(visible=False)), # Clear result and hide button first
inputs=None,
outputs=[result_image, use_as_input_button],
queue=False # Don't queue the clearing part
).then(
fn=infer, # Run the generation
inputs=gen_inputs,
outputs=result_image, # Output is the single generated image
).then(
fn=handle_output, # Process output: update history, show button
inputs=[result_image, history_gallery],
outputs=[result_image, history_gallery, use_as_input_button] # Update result again (no change), history, and button visibility
)
preview_button.click(
fn=preview_image_and_mask,
inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option, custom_resize_percentage, alignment_dropdown,
overlap_left, overlap_right, overlap_top, overlap_bottom],
outputs=preview_mask_image, # Output to the preview image component
queue=False # Preview should be fast
)
# Launch the app
demo.queue(max_size=12).launch(share=False, ssr_mode=False, show_error=True)