import gradio as gr import spaces import torch from diffusers import AutoencoderKL, TCDScheduler from diffusers.models.model_loading_utils import load_state_dict # Removed ImageSlider import from huggingface_hub import hf_hub_download # Ensure these custom modules are accessible in the environment # If running locally, they should be in the same directory or installed try: from controlnet_union import ControlNetModel_Union from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline except ImportError as e: print(f"Error importing custom modules: {e}") print("Please ensure 'controlnet_union.py' and 'pipeline_fill_sd_xl.py' are in the working directory or installed.") # Optionally, try installing if running in a suitable environment # import os # os.system("pip install git+https://github.com/UNION-AI-Research/FILL-Context-Aware-Outpainting.git") # Or wherever the package is hosted # Re-try import might be needed depending on environment setup exit() from PIL import Image, ImageDraw import numpy as np import os # For checking example files # --- Model Loading --- # Use environment variable for model cache if needed # HUGGINGFACE_HUB_CACHE = os.environ.get("HUGGINGFACE_HUB_CACHE", None) try: config_file = hf_hub_download( "xinsir/controlnet-union-sdxl-1.0", filename="config_promax.json", # cache_dir=HUGGINGFACE_HUB_CACHE ) 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", # cache_dir=HUGGINGFACE_HUB_CACHE ) sstate_dict = load_state_dict(model_file) model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model( controlnet_model, sstate_dict, model_file, "xinsir/controlnet-union-sdxl-1.0" ) model.to(device="cuda", dtype=torch.float16) print("ControlNet loaded successfully.") vae = AutoencoderKL.from_pretrained( "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, # cache_dir=HUGGINGFACE_HUB_CACHE ).to("cuda") print("VAE loaded successfully.") pipe = StableDiffusionXLFillPipeline.from_pretrained( "SG161222/RealVisXL_V5.0_Lightning", torch_dtype=torch.float16, vae=vae, controlnet=model, variant="fp16", # cache_dir=HUGGINGFACE_HUB_CACHE ).to("cuda") print("Pipeline loaded successfully.") pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config) print("Scheduler configured.") except Exception as e: print(f"Error during model loading: {e}") raise e # --- Helper Functions --- 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): if image is None: raise gr.Error("Input image not provided.") try: 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 elif resize_option == "Custom": resize_percentage = custom_resize_percentage else: raise ValueError(f"Invalid resize option: {resize_option}") # 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) # Ensure dimensions fit within target (can happen if original image is tiny and resize % is large) new_width = min(new_width, target_size[0]) new_height = min(new_height, target_size[1]) # 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 if overlap is enabled, otherwise 0 overlap_x = max(overlap_x, 1) if overlap_left or overlap_right else 0 overlap_y = max(overlap_y, 1) if overlap_top or overlap_bottom else 0 # 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 else: raise ValueError(f"Invalid alignment: {alignment}") # Adjust margins to ensure image is fully within bounds (should be redundant with min check above) 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)) # White background background.paste(source, (margin_x, margin_y)) # Create the mask (initially all black - meaning keep everything) mask_np = np.zeros(target_size[::-1], dtype=np.uint8) # Use numpy for easier slicing [::-1] for (height, width) # Calculate the coordinates of the *source image* area within the target canvas source_left = margin_x source_top = margin_y source_right = margin_x + new_width source_bottom = margin_y + new_height # Calculate the coordinates of the *unmasked* area (area to keep from source) unmasked_left = source_left + overlap_x if overlap_left else source_left unmasked_top = source_top + overlap_y if overlap_top else source_top unmasked_right = source_right - overlap_x if overlap_right else source_right unmasked_bottom = source_bottom - overlap_y if overlap_bottom else source_bottom # Special handling for edge alignments to ensure the edge itself is kept if overlap disabled if alignment == "Left" and not overlap_left: unmasked_left = source_left if alignment == "Right" and not overlap_right: unmasked_right = source_right if alignment == "Top" and not overlap_top: unmasked_top = source_top if alignment == "Bottom" and not overlap_bottom: unmasked_bottom = source_bottom # Ensure coordinates are valid and clipped to the source image area within the canvas unmasked_left = max(source_left, min(unmasked_left, source_right)) unmasked_top = max(source_top, min(unmasked_top, source_bottom)) unmasked_right = max(source_left, min(unmasked_right, source_right)) unmasked_bottom = max(source_top, min(unmasked_bottom, source_bottom)) # Create the final mask: White (255) = Area to inpaint/outpaint, Black (0) = Area to keep final_mask_np = np.ones(target_size[::-1], dtype=np.uint8) * 255 # Start with all white (change everything) if unmasked_right > unmasked_left and unmasked_bottom > unmasked_top: # Set the area to keep (calculated unmasked rectangle) to black (0) final_mask_np[unmasked_top:unmasked_bottom, unmasked_left:unmasked_right] = 0 mask = Image.fromarray(final_mask_np) return background, mask except Exception as e: print(f"Error in prepare_image_and_mask: {e}") raise gr.Error(f"Failed to prepare image and mask: {e}") def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom): if image is None: return None # Or return a placeholder image/message try: 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 for the masked (inpainting/outpainting) area red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 100)) # 100 alpha (~40% opacity) # The mask is white (255) where outpainting happens. Use this directly. preview.paste(red_overlay, (0, 0), mask) # Paste red where mask is white return preview except Exception as e: print(f"Error during preview generation: {e}") # Return the original background or an error placeholder if 'background' in locals(): return background.convert('RGBA') else: return Image.new('RGBA', (width, height), (200, 200, 200, 255)) # Grey placeholder @spaces.GPU(duration=60) # Adjusted duration slightly 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, progress=gr.Progress(track_tqdm=True)): if image is None: raise gr.Error("Please provide an input image.") try: # --- Preparation --- progress(0.1, desc="Preparing image and mask...") original_alignment = alignment 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) # --- Alignment Check & Correction --- # Get dimensions *after* initial placement and resize pasted_source_img_width = int(image.width * min(width / image.width, height / image.height) * (custom_resize_percentage if resize_option=='Custom' else {'Full':100, '50%':50, '33%':33, '25%':25}[resize_option])/100) pasted_source_img_height = int(image.height * min(width / image.width, height / image.height) * (custom_resize_percentage if resize_option=='Custom' else {'Full':100, '50%':50, '33%':33, '25%':25}[resize_option])/100) pasted_source_img_width = max(64, min(pasted_source_img_width, width)) pasted_source_img_height = max(64, min(pasted_source_img_height, height)) needs_reprepare = False if alignment in ("Left", "Right") and pasted_source_img_width >= width: print(f"Warning: Source width ({pasted_source_img_width}) >= target width ({width}) with {alignment} alignment. Forcing Middle alignment.") alignment = "Middle" needs_reprepare = True if alignment in ("Top", "Bottom") and pasted_source_img_height >= height: print(f"Warning: Source height ({pasted_source_img_height}) >= target height ({height}) with {alignment} alignment. Forcing Middle alignment.") alignment = "Middle" needs_reprepare = True if needs_reprepare and alignment != original_alignment: print("Re-preparing mask due to alignment change.") progress(0.15, desc="Re-preparing mask for Middle alignment...") 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) # ControlNet expects the image with the *original* content visible in the non-masked area cnet_image = background.copy() # In some ControlNet inpainting setups, you might mask the control image too, # but Union ControlNet Fill often works well with the unmasked source pasted onto the background. # cnet_image.paste(0, mask=ImageOps.invert(mask)) # Optional: Black out masked area in CNet image # --- Prompt Encoding --- progress(0.2, desc="Encoding prompt...") final_prompt = f"{prompt_input}, high quality, 4k" if prompt_input else "high quality, 4k" # Add default tags if no prompt negative_prompt = "low quality, blurry, noisy, text, words, letters, watermark, signature, username, artist name, deformed, distorted, disfigured, bad anatomy, extra limbs, missing limbs" # Note: TCD/Lightning pipelines often work better *without* explicit negative prompts encoded # Try encoding only the positive prompt first ( prompt_embeds, _, # negative_prompt_embeds (set to None or handle differently for TCD) pooled_prompt_embeds, _, # negative_pooled_prompt_embeds ) = pipe.encode_prompt(final_prompt, "cuda", False) # do_classifier_free_guidance=False for TCD # --- Inference --- progress(0.3, desc="Starting diffusion process...") print(f"Running inference with {num_inference_steps} steps...") pipeline_output = pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=None, # Pass None for TCD/Lightning pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=None, # Pass None for TCD/Lightning image=background, # Initial state for masked area (background with source) mask_image=mask, # Mask (white = change) control_image=cnet_image, # ControlNet input num_inference_steps=num_inference_steps, guidance_scale=0.0, # Crucial for TCD/Lightning controlnet_conditioning_scale=0.8, # Default for FILL pipeline, adjust if needed output_type="pil" # Ensure PIL output # Add tqdm=True if supported by the custom pipeline and using gr.Progress without track_tqdm ) # --- Process Output --- progress(0.9, desc="Processing results...") # Check if the pipeline returned a standard output object or a generator output_image = None if hasattr(pipeline_output, 'images'): # Standard diffusers output print("Pipeline returned a standard output object.") if len(pipeline_output.images) > 0: output_image = pipeline_output.images[0] else: raise ValueError("Pipeline output contained no images.") # Check if it's iterable (generator) - less likely with direct call and output_type='pil' but good practice elif hasattr(pipeline_output, '__iter__') and not isinstance(pipeline_output, dict): print("Pipeline returned a generator, iterating to get the final image.") last_item = None for item in pipeline_output: last_item = item # Try to extract image from the last yielded item (structure can vary) if isinstance(last_item, tuple) and len(last_item) > 0 and isinstance(last_item[0], Image.Image): output_image = last_item[0] elif isinstance(last_item, dict) and 'images' in last_item and len(last_item['images']) > 0: output_image = last_item['images'][0] elif isinstance(last_item, Image.Image): output_image = last_item elif hasattr(last_item, 'images') and len(last_item.images) > 0: # Handle case where object yielded early output_image = last_item.images[0] if output_image is None: raise ValueError("Pipeline generator did not yield a valid final image structure.") else: raise TypeError(f"Unexpected pipeline output type: {type(pipeline_output)}. Cannot extract image.") print("Inference complete.") progress(1.0, desc="Done!") return output_image except Exception as e: print(f"Error during inference: {e}") import traceback traceback.print_exc() # Print full traceback to console/logs raise gr.Error(f"Inference failed: {e}") def clear_result(*args): """Clears the result Image and related components.""" updates = { result: gr.update(value=None), use_as_input_button: gr.update(visible=False), } # If preview image is passed as an arg, clear it too if len(args) > 0 and isinstance(args[0], gr.Image): updates[args[0]] = gr.update(value=None) # Assuming preview_image is the first optional arg return updates # --- UI Helper Functions --- def preload_presets(target_ratio, ui_width, ui_height): """Updates the width and height sliders based on the selected aspect ratio.""" settings_update = gr.update() # Default: no change to accordion state if target_ratio == "9:16": changed_width = 720 changed_height = 1280 elif target_ratio == "16:9": changed_width = 1280 changed_height = 720 elif target_ratio == "1:1": changed_width = 1024 changed_height = 1024 elif target_ratio == "Custom": changed_width = ui_width # Keep current slider values changed_height = ui_height settings_update = gr.update(open=True) # Open accordion for custom else: # Should not happen changed_width = ui_width changed_height = ui_height return changed_width, changed_height, settings_update def select_the_right_preset(user_width, user_height): """Updates the radio button based on the current slider values.""" 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 not isinstance(new_image, Image.Image): # Don't add if generation failed (None) return history or [] # Return current or empty list if history is None: history = [] history.insert(0, new_image) # Limit history size (optional) max_history = 12 if len(history) > max_history: history = history[:max_history] return history # --- Gradio UI Definition --- css = """ .gradio-container { max-width: 1200px !important; /* Use max-width for responsiveness */ margin: auto !important; /* Center the container */ padding: 10px; /* Add some padding */ } h1 { text-align: center; margin-bottom: 15px;} footer { display: none !important; /* More reliable way to hide footer */ } /* Ensure result image takes reasonable space */ #result-image img { max-height: 768px; /* Adjust max height as needed */ object-fit: contain; width: 100%; /* Allow image to use column width */ height: auto; display: block; /* Prevent extra space below image */ margin: auto; /* Center image within its container */ } #input-image img { max-height: 400px; object-fit: contain; width: 100%; height: auto; display: block; margin: auto; } #preview-image img { max-height: 250px; /* Smaller preview */ object-fit: contain; width: 100%; height: auto; display: block; margin: auto; } #history-gallery .thumbnail-item { /* Style history items */ height: 100px !important; overflow: hidden; /* Hide overflow */ } #history-gallery .gallery { grid-template-rows: repeat(auto-fill, 100px) !important; gap: 4px !important; /* Add small gap */ } #history-gallery .thumbnail-item img { object-fit: contain !important; /* Ensure history previews fit */ height: 100%; width: 100%; } /* Make Checkboxes smaller and closer */ .gradio-checkboxgroup .wrap { gap: 0.5rem 1rem !important; /* Adjust spacing */ } .gradio-checkbox label span { font-size: 0.9em; /* Slightly smaller label text */ } .gradio-checkbox input { transform: scale(0.9); /* Slightly smaller checkbox */ } /* Style Accordion */ .gradio-accordion .label-wrap { /* Target the label wrapper */ border: 1px solid #e0e0e0; border-radius: 5px; padding: 8px 12px; background-color: #f9f9f9; } """ title = """