pipelines.py

"""
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