""" Contains functions to execute the main image generation stages: 1. OpenPose Detection: Extracts pose information. 2. Low-Resolution Generation: Creates initial image using Pose ControlNet. 3. High-Resolution Tiling: Upscales the low-res image using Tile ControlNet. Manages dynamic loading/unloading of diffusion pipelines to conserve VRAM. """ import torch import gc import time import os from PIL import Image from tqdm.auto import tqdm import gradio as gr from diffusers import ( StableDiffusionControlNetImg2ImgPipeline, UniPCMultistepScheduler, ) from model_loader import ( get_openpose_detector, get_controlnet_pose, get_controlnet_tile, get_device, get_dtype, are_models_loaded, ) from image_utils import create_blend_mask from prompts import get_prompts_for_run # --- Configuration --- BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5" LORA_DIR = "loras" LORA_FILES = { "style": os.path.join(LORA_DIR, "night_comic_V06.safetensors"), "detail": os.path.join(LORA_DIR, "add_detail.safetensors"), } LORA_WEIGHTS_LOWRES = [1, 1] LORA_WEIGHTS_HIRES = [1, 2] ACTIVE_ADAPTERS = ["style", "detail"] def cleanup_memory(): """Forces garbage collection and clears CUDA cache.""" gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # --- Stage 1: OpenPose Detection --- def run_pose_detection(resized_input_image): """ Detects human pose (body, hands, face) from the input image using OpenPose. Temporarily moves the OpenPose detector model to the active GPU (if available) for processing and then moves it back to the CPU to conserve VRAM. Args: input_image_resized (PIL.Image.Image): The input image, already resized and in RGB format. Returns: PIL.Image.Image | None: A PIL Image representing the detected pose map, or None if detection fails or models aren't loaded. """ if not are_models_loaded(): print("Error: Cannot run pose detection, models not loaded.") return None detector = get_openpose_detector() device = get_device() control_image_openpose = None if detector is None: print("Error: OpenPose detector is None.") return None try: detector.to(device) cleanup_memory() control_image_openpose = detector( resized_input_image, include_face=True, include_hand=True ) except Exception as e: print(f"ERROR during OpenPose detection: {e}") control_image_openpose = None finally: detector.to("cpu") cleanup_memory() return control_image_openpose # --- Stage 2: Low-Resolution Generation --- def run_low_res_generation( resized_input_image, pose_map, seed, steps, guidance_scale, strength, controlnet_scale=0.8, progress=gr.Progress(track_tqdm=True) ): """ Generates the initial low-resolution image using Img2Img with Pose ControlNet. Dynamically loads the StableDiffusionControlNetImg2ImgPipeline, applies LoRAs, runs inference, and then unloads the pipeline to free VRAM before returning. Args: input_image_resized (PIL.Image.Image): The resized input image. pose_map (PIL.Image.Image): The pose map generated by run_pose_detection. seed (int): The random seed for generation. steps (int): Number of diffusion inference steps. guidance_scale (float): Classifier-free guidance scale. strength (float): Img2Img strength (0.0 to 1.0). How much noise to add. controlnet_scale (float): Conditioning scale for the Pose ControlNet. progress (gr.Progress): Gradio progress object for UI updates. Returns: PIL.Image.Image | None: The generated low-resolution PIL Image, or None if an error occurs. Raises: gr.Error: Raises a Gradio error if generation fails catastrophically. """ if not are_models_loaded() or pose_map is None: error_msg = "Cannot run low-res generation: " if not are_models_loaded(): error_msg += "Models not loaded. " if pose_map is None: error_msg += "Pose map is missing." print(f"Error: {error_msg}") return None device = get_device() dtype = get_dtype() controlnet_pose = get_controlnet_pose() output_image_lowres = None pipe_lowres = None positive_prompt, negative_prompt, _, _ = get_prompts_for_run() generator = torch.Generator(device=device).manual_seed(int(seed)) progress(0, desc="Loading Low-Res Pipeline...") try: # 1. Load Pipeline pipe_lowres = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( BASE_MODEL_ID, controlnet=controlnet_pose, torch_dtype=dtype, safety_checker=None ) pipe_lowres.scheduler = UniPCMultistepScheduler.from_config(pipe_lowres.scheduler.config) pipe_lowres.to(device) cleanup_memory() # 2. Load LoRAs if os.path.exists(LORA_FILES["style"]) and os.path.exists(LORA_FILES["detail"]): pipe_lowres.load_lora_weights(LORA_FILES["style"], adapter_name="style") pipe_lowres.load_lora_weights(LORA_FILES["detail"], adapter_name="detail") pipe_lowres.set_adapters(ACTIVE_ADAPTERS, adapter_weights=LORA_WEIGHTS_LOWRES) print(f"Activated LoRAs: {ACTIVE_ADAPTERS} with weights {LORA_WEIGHTS_LOWRES}") else: print("Warning: One or both LoRA files not found. Skipping LoRA loading.") raise gr.Error("Required LoRA files not found in loras/ directory.") # 3. Run Inference progress(0.3, desc="Generating Low-Res Image...") output_image_low_res = pipe_lowres( prompt=positive_prompt, negative_prompt=negative_prompt, image=resized_input_image, control_image=pose_map, num_inference_steps=int(steps), strength=strength, guidance_scale=guidance_scale, controlnet_conditioning_scale=float(controlnet_scale), generator=generator, ).images[0] progress(0.9, desc="Low-Res Complete") except Exception as e: print(f"ERROR during Low-Res Generation Pipeline: {e}") import traceback traceback.print_exc() output_image_low_res = None raise gr.Error(f"Failed during low-res generation: {e}") finally: # 4. Cleanup Pipeline print("Cleaning up Low-Res pipeline...") if pipe_lowres is not None: try: if hasattr(pipe_lowres, 'get_active_adapters') and pipe_lowres.get_active_adapters(): print("Unloading LoRAs...") pipe_lowres.unload_lora_weights() except Exception as unload_e: print(f"Note: Error unloading LoRAs: {unload_e}") print("Moving Low-Res pipe components to CPU before deleting...") try: pipe_lowres.to('cpu') except Exception as cpu_e: print(f"Note: Error moving pipe to CPU: {cpu_e}") print("Deleting Low-Res pipeline object...") del pipe_lowres pipe_lowres = None print("Running garbage collection and emptying CUDA cache after Low-Res...") cleanup_memory() # time.sleep(1) print("--- Low-Res Generation Stage Finished ---") return output_image_low_res # --- Stage 3: High-Resolution Tiling Upscaling --- def run_hires_tiling( low_res_image, seed, steps, guidance_scale, strength, controlnet_scale=1.0, upscale_factor=2, tile_size=1024, tile_stride=1024, progress=gr.Progress(track_tqdm=True) ): """ Upscales the low-resolution image using tiling with the Tile ControlNet. Dynamically loads the StableDiffusionControlNetImg2ImgPipeline for tiling, applies LoRAs, processes the image in overlapping tiles, blends the results, and unloads the pipeline to free VRAM. Args: low_res_image (PIL.Image.Image): The low-resolution image from the previous stage. seed (int): The random seed (should ideally match low-res stage seed). steps (int): Number of diffusion inference steps per tile. guidance_scale (float): Classifier-free guidance scale for tiles. strength (float): Img2Img strength for tiling (controls detail vs. original). controlnet_scale (float): Conditioning scale for the Tile ControlNet. upscale_factor (int): Factor by which to increase the image resolution. tile_size (int): Size of the square tiles to process. tile_stride (int): Step size between tiles. Overlap = tile_size - tile_stride. progress (gr.Progress): Gradio progress object for UI updates. Returns: PIL.Image.Image | None: The generated high-resolution PIL Image, or None if an error occurs. Raises: gr.Error: Raises a Gradio error if tiling fails catastrophically. """ if not are_models_loaded() or low_res_image is None: error_msg = "Cannot run hi-res tiling: " if not are_models_loaded(): error_msg += "Models not loaded. " if low_res_image is None: error_msg += "Low-res image is missing." print(f"Error: {error_msg}") return None device = get_device() dtype = get_dtype() controlnet_tile = get_controlnet_tile() high_res_output_image = None pipe_hires = None _, _, positive_prompt_tile, negative_prompt_tile = get_prompts_for_run() generator_tile = torch.Generator(device=device).manual_seed(int(seed)) print("\n--- Starting Hi-Res Tiling Stage ---") progress(0, desc="Preparing for Tiling...") try: # --- Setup Tiling Parameters --- target_width = low_res_image.width * upscale_factor target_height = low_res_image.height * upscale_factor if tile_size > min(target_width, target_height): print(f"Warning: Tile size ({tile_size}) > target dimension ({target_width}x{target_height}). Clamping tile size.") tile_size = min(target_width, target_height) tile_stride = tile_size overlap = tile_size - tile_stride if overlap < 0: print("Warning: Tile stride is larger than tile size. Setting stride = tile size.") tile_stride = tile_size overlap = 0 print(f"Target Res: {target_width}x{target_height}, Tile Size: {tile_size}, Stride: {tile_stride}, Overlap: {overlap}") # 1. Load Pipeline print(f"Loading Hi-Res Pipeline ({BASE_MODEL_ID} + Tile ControlNet)...") progress(0.05, desc="Loading Hi-Res Pipeline...") pipe_hires = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( BASE_MODEL_ID, controlnet=controlnet_tile, torch_dtype=dtype, safety_checker=None, ) pipe_hires.scheduler = UniPCMultistepScheduler.from_config(pipe_hires.scheduler.config) pipe_hires.to(device) # pipe_hires.enable_model_cpu_offload() # pipe_hires.enable_xformers_memory_efficient_attention() print("Hi-Res Pipeline loaded to GPU.") cleanup_memory() # 2. Load LoRAs print("Loading LoRAs for Hi-Res pipe...") if os.path.exists(LORA_FILES["style"]) and os.path.exists(LORA_FILES["detail"]): pipe_hires.load_lora_weights(LORA_FILES["style"], adapter_name="style") pipe_hires.load_lora_weights(LORA_FILES["detail"], adapter_name="detail") pipe_hires.set_adapters(ACTIVE_ADAPTERS, adapter_weights=LORA_WEIGHTS_HIRES) print(f"Activated LoRAs: {ACTIVE_ADAPTERS} with weights {LORA_WEIGHTS_HIRES}") else: print("Warning: One or both LoRA files not found. Skipping LoRA loading.") raise gr.Error("Required LoRA files not found in loras/ directory.") # --- Prepare for Tiling Loop --- print(f"Creating blurry base image ({target_width}x{target_height})...") progress(0.15, desc="Preparing Base Image...") blurry_high_res = low_res_image.resize((target_width, target_height), Image.LANCZOS) final_image = Image.new("RGB", (target_width, target_height)) blend_mask = create_blend_mask(tile_size, overlap) num_tiles_x = (target_width + tile_stride - 1) // tile_stride num_tiles_y = (target_height + tile_stride - 1) // tile_stride total_tiles = num_tiles_x * num_tiles_y print(f"Processing {num_tiles_x}x{num_tiles_y} = {total_tiles} tiles...") # --- Tiling Loop --- progress(0.2, desc=f"Processing Tiles (0/{total_tiles})") processed_tile_count = 0 with tqdm(total=total_tiles, desc="Tiling Upscale") as pbar: for y in range(num_tiles_y): for x in range(num_tiles_x): tile_start_time = time.time() pbar.set_description(f"Tiling Upscale (Tile {processed_tile_count+1}/{total_tiles})") x_start = x * tile_stride y_start = y * tile_stride x_end = min(x_start + tile_size, target_width) y_end = min(y_start + tile_size, target_height) crop_box = (x_start, y_start, x_end, y_end) tile_image_blurry = blurry_high_res.crop(crop_box) current_tile_width, current_tile_height = tile_image_blurry.size if current_tile_width < tile_size or current_tile_height < tile_size: try: edge_color = tile_image_blurry.getpixel((0, 0)) except IndexError: edge_color = (127, 127, 127) padded_tile = Image.new("RGB", (tile_size, tile_size), edge_color) padded_tile.paste(tile_image_blurry, (0, 0)) tile_image_blurry = padded_tile print(f"Padded edge tile at ({x},{y})") # 3. Run Inference on the Tile with torch.inference_mode(): output_tile = pipe_hires( prompt=positive_prompt_tile, negative_prompt=negative_prompt_tile, image=tile_image_blurry, control_image=tile_image_blurry, num_inference_steps=int(steps), strength=strength, guidance_scale=guidance_scale, controlnet_conditioning_scale=float(controlnet_scale), generator=generator_tile, output_type="pil" ).images[0] # --- Stitch Tile Back --- paste_x = x_start paste_y = y_start crop_w = x_end - x_start crop_h = y_end - y_start output_tile_region = output_tile.crop((0, 0, crop_w, crop_h)) if overlap > 0: blend_mask_region = blend_mask.crop((0, 0, crop_w, crop_h)) current_content_region = final_image.crop((paste_x, paste_y, paste_x + crop_w, paste_y + crop_h)) blended_tile_region = Image.composite(output_tile_region, current_content_region, blend_mask_region) final_image.paste(blended_tile_region, (paste_x, paste_y)) else: final_image.paste(output_tile_region, (paste_x, paste_y)) processed_tile_count += 1 pbar.update(1) # Update Gradio progress gradio_progress = 0.2 + 0.75 * (processed_tile_count / total_tiles) progress(gradio_progress, desc=f"Processing Tile {processed_tile_count}/{total_tiles}") tile_end_time = time.time() print(f"Tile ({x},{y}) processed in {tile_end_time - tile_start_time:.2f}s") # cleanup_memory() print("Tile processing complete.") high_res_output_image = final_image progress(0.95, desc="Tiling Complete") except Exception as e: print(f"ERROR during Hi-Res Tiling Pipeline: {e}") import traceback traceback.print_exc() high_res_output_image = None raise gr.Error(f"Failed during hi-res tiling: {e}") finally: # 4. Cleanup Pipeline print("Cleaning up Hi-Res pipeline...") if pipe_hires is not None: try: if hasattr(pipe_hires, 'get_active_adapters') and pipe_hires.get_active_adapters(): print("Unloading LoRAs...") pipe_hires.unload_lora_weights() except Exception as unload_e: print(f"Note: Error unloading LoRAs: {unload_e}") print("Moving Hi-Res pipe components to CPU before deleting...") try: pipe_hires.to('cpu') except Exception as cpu_e: print(f"Note: Error moving pipe to CPU: {cpu_e}") print("Deleting Hi-Res pipeline object...") del pipe_hires pipe_hires = None print("Running garbage collection and emptying CUDA cache after Hi-Res...") cleanup_memory() print("--- Hi-Res Tiling Stage Finished ---") return high_res_output_image