BrushNet

Implements a specialized dual-branch architecture that separates masked image features from noisy latent features during the diffusion process, reducing the model's learning load and enabling precise inpainting. The architecture processes segmentation or random masks through dedicated branches that converge at multiple resolution levels, allowing the base diffusion model to focus on content generation within masked regions while preserving unmasked areas. This decomposition is achieved through custom UNet modifications in the diffusers library that inject BrushNet control at intermediate layers without requiring full model retraining.

Provides unified inference pipelines (StableDiffusionBrushNetPipeline and StableDiffusionXLBrushNetPipeline) that orchestrate the complete inpainting workflow: text encoding via CLIP/OpenCLIP, mask preprocessing, latent encoding of the original image, iterative diffusion with BrushNet control injection, and final decoding. The pipeline abstracts away the complexity of managing multiple model components (text encoder, VAE, UNet, scheduler) and provides a simple API while supporting both SD 1.5 and SDXL base models with separate segmentation and random mask variants.

Provides tools for reducing model size and inference latency through quantization (INT8, FP16) and optimization techniques. The system supports post-training quantization of BrushNet weights, mixed-precision inference (FP16 for forward pass, FP32 for critical operations), and optional pruning of less important weights. Quantized models achieve 2-4x speedup with minimal quality loss, enabling deployment on resource-constrained devices (edge GPUs, mobile) or higher throughput on servers.

Provides seamless integration with the HuggingFace diffusers library, enabling BrushNet to work with any diffusers-compatible scheduler, pipeline, and model. The integration includes custom BrushNet model classes (BrushNetModel) that inherit from diffusers base classes, custom pipeline classes (StableDiffusionBrushNetPipeline) that follow diffusers conventions, and compatibility with diffusers utilities (safety checker, feature extractor). This enables users to leverage the entire diffusers ecosystem (LoRA, ControlNet, other extensions) alongside BrushNet.

Preprocesses input images and masks into latent space representations that preserve spatial information about masked vs unmasked regions. The system encodes the original image through the VAE encoder, then applies mask-aware feature extraction that separates masked image features from the noisy latent representation. This preprocessing step is critical for the dual-branch architecture, as it ensures the BrushNet model receives properly formatted input that distinguishes between regions to inpaint and regions to preserve, using spatial masking operations at the latent level (typically 8x downsampled from image space).

Injects BrushNet control signals at multiple UNet resolution levels (typically 4 scales: 64x64, 32x32, 16x16, 8x8) to provide fine-grained guidance over the diffusion process. The control mechanism works by modifying the UNet's cross-attention and self-attention layers with BrushNet-specific conditioning that incorporates mask information and masked image features at each resolution. This multi-scale injection ensures that both coarse structure (from low-resolution features) and fine details (from high-resolution features) are properly controlled, enabling precise inpainting without affecting unmasked regions.

Provides separate model variants optimized for two distinct mask types: segmentation masks (clean, object-shaped boundaries) and random masks (arbitrary, potentially irregular shapes). Each variant is trained with different mask distributions and augmentation strategies to handle the specific characteristics of its target mask type. The system automatically selects the appropriate variant based on mask properties or allows explicit selection, enabling optimal inpainting quality for different use cases without requiring users to understand the underlying mask type differences.

Provides end-to-end training infrastructure for fine-tuning BrushNet on custom datasets, including dataset loading, mask generation/augmentation, and training loop management. The training system supports both SD 1.5 and SDXL base models with separate training scripts, implements mask augmentation strategies (random mask generation, boundary noise, dilation/erosion), and uses mixed-precision training with gradient accumulation for memory efficiency. Training can be performed on standard datasets (Places, CelebA-HQ) or custom image collections, with support for distributed training across multiple GPUs.

Computes standard image quality metrics for evaluating inpainting results: LPIPS (learned perceptual image patch similarity) for perceptual quality, FID (Fréchet Inception Distance) for distribution matching, and SSIM (structural similarity) for pixel-level fidelity. The evaluation system loads pre-trained feature extractors (InceptionV3 for FID, AlexNet for LPIPS) and compares generated inpainted images against ground truth or reference images. Results are aggregated across test sets and reported with statistical summaries (mean, std, percentiles).

Provides a browser-based interactive interface for real-time inpainting using Gradio, enabling users to upload images, draw masks, enter text prompts, and adjust inference parameters (guidance scale, steps) without coding. The interface handles image upload, mask drawing with canvas tools, prompt input, and displays results with latency information. The Gradio app wraps the inference pipeline and can be deployed locally or on cloud platforms (HuggingFace Spaces, Gradio Cloud) for easy sharing and collaboration.

Extends basic text-guided inpainting with instruction-based editing that interprets natural language instructions to automatically generate masks and guide inpainting. The system parses instructions like 'remove the person on the left' or 'replace the sky with clouds' to identify regions of interest and apply appropriate inpainting. This capability combines text understanding with spatial reasoning, potentially using auxiliary models (object detection, segmentation) to convert instructions into masks before applying BrushNet inpainting.

Enables efficient processing of multiple images with different masks and prompts in a single batch, optimizing GPU utilization and reducing per-image overhead. The batch processor handles variable image sizes through padding/resizing, manages memory efficiently with dynamic batching, and provides progress tracking and error handling for robust production use. Results are returned with metadata (processing time, success/failure status) for each image.