app.py

# -*- coding: utf-8 -*-
"""controlnet_depth_canny_segmentation.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1oeKagV0PyeA1ezzMP4h4KnRRHaklw7L1
"""

import os
import torch
from PIL import Image
import numpy as np
import cv2
import random
import gradio as gr

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from transformers import AutoTokenizer, CLIPTextModel, CLIPFeatureExtractor
from transformers import DPTForDepthEstimation, DPTImageProcessor
# Imports for Segmentation
from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor

stable_diffusion_base = "runwayml/stable-diffusion-v1-5"

# Path to your fine-tuned ControlNet Depth model
finetune_controlnet_depth_path = "controlnet"

# Pre-trained ControlNet models for Canny and Segmentation
# These are standard models from the lllyasviel collection, optimized for these tasks.
controlnet_canny_pretrained_path = "lllyasviel/sd-controlnet-canny"
controlnet_seg_pretrained_path = "lllyasviel/sd-controlnet-seg"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32

# Global pipeline and pre-processor models
pipeline = None
depth_estimator_model = None
depth_estimator_processor = None
segmentation_model_preprocessor = None # Renamed to avoid conflict with ControlNet model
segmentation_processor_preprocessor = None # Renamed to avoid conflict with ControlNet model

# Global ControlNet models
controlnet_depth_model = None
controlnet_canny_model = None
controlnet_seg_model = None

def load_depth_estimator():
    """Loads the MiDaS depth estimation model."""
    global depth_estimator_model, depth_estimator_processor

    if depth_estimator_model is None:

        print("Loading LiheYoung/depth-anything-large-hf depth estimation model...")
        model_name = "LiheYoung/depth-anything-large-hf"

        depth_estimator_model = DPTForDepthEstimation.from_pretrained(model_name)
        depth_estimator_processor = DPTImageProcessor.from_pretrained(model_name)
        depth_estimator_model.to(DEVICE)

        depth_estimator_model.eval()
        print("MiDaS depth estimation model loaded.")


    return depth_estimator_model, depth_estimator_processor

def load_segmentation_preprocessor_model():
    """Loads the Mask2Former segmentation pre-processor model."""
    global segmentation_model_preprocessor, segmentation_processor_preprocessor

    if segmentation_model_preprocessor is None:
        print("Loading Mask2Former segmentation pre-processor model...")
        model_name = "facebook/mask2former-swin-large-ade-semantic"

        segmentation_processor_preprocessor = Mask2FormerImageProcessor.from_pretrained(model_name)
        segmentation_model_preprocessor = Mask2FormerForUniversalSegmentation.from_pretrained(model_name)
        segmentation_model_preprocessor.to(DEVICE)

        segmentation_model_preprocessor.eval()
        print("Mask2Former segmentation pre-processor model loaded.")
    return segmentation_model_preprocessor, segmentation_processor_preprocessor

def load_diffusion_pipeline_and_controlnets():
    """
    Loads the base Stable Diffusion pipeline components and all ControlNet models.
    """
    global pipeline, controlnet_depth_model, controlnet_canny_model, controlnet_seg_model

    if pipeline is None:
        print("Loading base Stable Diffusion pipeline components...")

        try:
            # Load individual components of the base Stable Diffusion pipeline
            vae = AutoencoderKL.from_pretrained(stable_diffusion_base, subfolder="vae", torch_dtype=DTYPE)
            tokenizer = AutoTokenizer.from_pretrained(stable_diffusion_base, subfolder="tokenizer")
            text_encoder = CLIPTextModel.from_pretrained(stable_diffusion_base, subfolder="text_encoder", torch_dtype=DTYPE)
            unet = UNet2DConditionModel.from_pretrained(stable_diffusion_base, subfolder="unet", torch_dtype=DTYPE)
            scheduler = DDPMScheduler.from_pretrained(stable_diffusion_base, subfolder="scheduler")
            feature_extractor = CLIPFeatureExtractor.from_pretrained(stable_diffusion_base, subfolder="feature_extractor")

            # Load ControlNet models
            print("Loading ControlNet models (Depth, Canny, Segmentation)...")
            if not os.path.exists(finetune_controlnet_depth_path):
                raise FileNotFoundError(f"Fine-tuned ControlNet Depth model not found at: {finetune_controlnet_depth_path}")

            controlnet_depth_model = ControlNetModel.from_pretrained(finetune_controlnet_depth_path, torch_dtype=DTYPE)
            controlnet_canny_model = ControlNetModel.from_pretrained(controlnet_canny_pretrained_path, torch_dtype=DTYPE)
            controlnet_seg_model = ControlNetModel.from_pretrained(controlnet_seg_pretrained_path, torch_dtype=DTYPE)
            print("All ControlNet models loaded.")

            # Create the StableDiffusionControlNetPipeline with a list of ControlNets
            pipeline = StableDiffusionControlNetPipeline(
                vae=vae,
                text_encoder=text_encoder,
                tokenizer=tokenizer,
                unet=unet,
                controlnet=[controlnet_depth_model, controlnet_canny_model, controlnet_seg_model], # Pass list of ControlNets
                scheduler=scheduler,
                safety_checker=None,
                feature_extractor=feature_extractor,
                image_encoder=None,
                requires_safety_checker=False,
            )

            pipeline.to(DEVICE)
            if torch.cuda.is_available() and hasattr(pipeline, "enable_xformers_memory_efficient_attention"):
                try:
                    pipeline.enable_xformers_memory_efficient_attention()
                    print("xformers memory efficient attention enabled.")
                except Exception as e:
                    print(f"Could not enable xformers: {e}")

            # Load all necessary pre-processor models at startup
            load_depth_estimator()
            load_segmentation_preprocessor_model()
            print("Diffusion pipeline and all pre-processor models loaded successfully.")

        except Exception as e:
            print(f"Error loading pipeline or ControlNets: {e}")
            pipeline = None
            raise RuntimeError(f"Failed to load diffusion pipeline or ControlNets: {e}")
    return pipeline

def estimate_depth(pil_image: Image.Image) -> Image.Image:
    """Estimates depth map from a PIL image."""
    global depth_estimator_model, depth_estimator_processor

    if depth_estimator_model is None or depth_estimator_processor is None:
        try:
            load_depth_estimator()
        except RuntimeError as e:
            raise RuntimeError(f"Depth estimator not loaded: {e}")

    inputs = depth_estimator_processor(images=pil_image, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = depth_estimator_model(**inputs)
        predicted_depth = outputs.predicted_depth

    depth_numpy = predicted_depth.squeeze().cpu().numpy()

    min_depth = depth_numpy.min()
    max_depth = depth_numpy.max()
    normalized_depth = (depth_numpy - min_depth) / (max_depth - min_depth)

    inverted_normalized_depth = 1 - normalized_depth

    depth_image_array = (inverted_normalized_depth * 255).astype(np.uint8)
    depth_pil_image = Image.fromarray(depth_image_array).convert("RGB")
    return depth_pil_image

def estimate_canny(pil_image: Image.Image) -> Image.Image:
    """Estimates Canny edges from a PIL image."""
    img = np.array(pil_image)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    low_threshold = 100
    high_threshold = 200
    canny = cv2.Canny(blurred, low_threshold, high_threshold)
    canny_pil = Image.fromarray(canny).convert("RGB")
    return canny_pil

def estimate_segmentation(pil_image: Image.Image) -> Image.Image:
    """Estimates segmentation map from a PIL image."""
    global segmentation_model_preprocessor, segmentation_processor_preprocessor
    if segmentation_model_preprocessor is None or segmentation_processor_preprocessor is None:
        try:
            load_segmentation_preprocessor_model()
        except RuntimeError as e:
            raise RuntimeError(f"Segmentation model not loaded: {e}")

    inputs = segmentation_processor_preprocessor(images=pil_image, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = segmentation_model_preprocessor(**inputs)

    # The Mask2FormerForUniversalSegmentation model returns 'sem_seg' for semantic segmentation
    # I will use the post_process_semantic_segmentation function to get the processed mask
    target_size = pil_image.size[::-1] # (height, width)
    segmentation_maps = segmentation_processor_preprocessor.post_process_semantic_segmentation(
        outputs, target_sizes=[target_size]
    )
    # Access the processed segmentation map from the list
    predicted_mask = segmentation_maps[0].squeeze(0).cpu().numpy()


    unique_classes = np.unique(predicted_mask)
    color_map = {}
    for class_id in unique_classes:
        if class_id == 0:
            color_map[class_id] = (0, 0, 0) # Black for background
        else:
            # Generate consistent colors for classes across runs, but still random-ish
            # Convert class_id to standard integer for random.seed()
            random.seed(int(class_id))
            color_map[class_id] = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))

    colored_segmentation_array = np.zeros((*predicted_mask.shape, 3), dtype=np.uint8)
    for y in range(predicted_mask.shape[0]):
        for x in range(predicted_mask.shape[1]):
            colored_segmentation_array[y, x] = color_map[predicted_mask[y, x]]

    segmentation_pil_image = Image.fromarray(colored_segmentation_array).convert("RGB")
    return segmentation_pil_image

def generate_image_with_all_controls_simultaneous(
    input_image_raw: Image.Image,
    prompt: str,
    negative_prompt: str = "",
    num_inference_steps: int = 25,
    guidance_scale: float = 8.0,
    strength: float = 0.8,
    seed: int = None,
    resolution: int = 512
) -> tuple[Image.Image, Image.Image, Image.Image, Image.Image]: # Returns generated image + 3 control maps

    global pipeline
    if pipeline is None:
        try:
            load_diffusion_pipeline_and_controlnets()
        except RuntimeError as e:
            # Raise a Gradio Error instead of returning it
            raise gr.Error(f"Model not loaded: {e}")

    # 1. Generate all control maps
    print("Generating all control maps (Depth, Canny, Segmentation)...")
    try:
        depth_map_pil = estimate_depth(input_image_raw)
        canny_map_pil = estimate_canny(input_image_raw)
        segmentation_map_pil = estimate_segmentation(input_image_raw)
        print("All control maps generated.")
    except Exception as e:
        # Raise a Gradio Error instead of returning it
        raise gr.Error(f"Error during control map generation: {e}")

    print(f"Generating image for prompt: '{prompt}' (Negative: '{negative_prompt}', Strength: {strength})")

    # Resize all generated control maps to the desired resolution
    # IMPORTANT: The order here must match the order of ControlNet models in the pipeline
    # (depth_model, canny_model, seg_model)
    control_images_for_pipeline = [
        depth_map_pil.resize((resolution, resolution), Image.LANCZOS).convert("RGB"),
        canny_map_pil.resize((resolution, resolution), Image.LANCZOS).convert("RGB"),
        segmentation_map_pil.resize((resolution, resolution), Image.LANCZOS).convert("RGB")
    ]

    generator = None
    if seed is None:
        seed = random.randint(0, 100000)
    generator = torch.Generator(device=DEVICE).manual_seed(seed)

    with torch.no_grad():
        generated_images = pipeline(
            prompt,
            negative_prompt=negative_prompt,
            image=control_images_for_pipeline, # Pass the list of control images
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            strength=strength,
            generator=generator,
        ).images

    print(f"Image generation complete (seed: {seed}).")
    # Return the generated image and all three control maps for visualization
    return generated_images[0], depth_map_pil, canny_map_pil, segmentation_map_pil

# Gradio Interface Setup
with gr.Blocks() as iface:
    gr.Markdown(
        """
        # Stable Diffusion ControlNet Multi-Control (Simultaneous) Demo
        Upload an input image, and the app will generate its **Depth Map**, **Canny Edges**, and **Segmentation Map**.
        These three control maps will then be used **simultaneously** with your text prompt to generate a new image.
        This provides highly detailed structural guidance.

        **⚠️ WARNING: This setup requires significant GPU memory. It may crash on smaller GPUs (e.g., Colab T4).**
        """
    )

    with gr.Row():
        with gr.Column():
            input_image_raw = gr.Image(type="pil", label="Input Image")
            prompt = gr.Textbox(label="Prompt", value="a high-quality photo of a modern interior design, photorealistic, 4k")
            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, low quality, bad anatomy, deformed, ugly, disfigured, watermark, text, signature, error, missing limbs, extra limbs, mutated, out of frame, cropped, noisy, grainy, jpeg artifacts, cartoon, painting, illustration, sketch, drawing, 3d render", placeholder="Enter negative prompt to guide generation away from these features")
            num_inference_steps = gr.Slider(minimum=10, maximum=100, value=25, step=1, label="Inference Steps")
            guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=8.0, step=0.5, label="Guidance Scale")
            strength = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.01, label="Strength (0.0-1.0)")
            seed = gr.Number(label="Seed (optional, leave blank for random)", value=None)
            resolution = gr.Number(label="Resolution", value=512, interactive=False)
            submit_btn = gr.Button("Generate Images")

        with gr.Column():
            generated_image_output = gr.Image(type="pil", label="Generated Image (Multi-Control)")
            with gr.Row():
                depth_map_output = gr.Image(type="pil", label="Generated Depth Map")
                canny_map_output = gr.Image(type="pil", label="Generated Canny Edges")
                segmentation_map_output = gr.Image(type="pil", label="Generated Segmentation Map")

    # Define the action for the submit button
    submit_btn.click(
        fn=generate_image_with_all_controls_simultaneous,
        inputs=[
            input_image_raw,
            prompt,
            negative_prompt,
            num_inference_steps,
            guidance_scale,
            strength,
            seed,
            resolution
        ],
        outputs=[
            generated_image_output,
            depth_map_output,
            canny_map_output,
            segmentation_map_output
        ]
    )

# Load the pipeline and pre-processor models when the Gradio app starts
load_diffusion_pipeline_and_controlnets()

if __name__ == "__main__":
    iface.launch(debug=True, share=True)