# -*- 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)