Upload controlnet_depth_canny_segmentation.py

controlnet_depth_canny_segmentation.py

@@ -0,0 +1,351 @@
+# -*- 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 = "/content/drive/MyDrive/Output_ControlNet_Finetune"
+
+# 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)
+
+from google.colab import drive
+drive.mount('/content/drive')