ControlNet Implementation in PyTorch

This implementation builds a minimal ControlNet from scratch. It includes a ZeroConv2d that initializes a Conv2d’s weights and bias to zero, a ControlNet module with a control encoder made of stacked conv-BN-ReLU blocks, residual skip connections, channel attention via a SE-style squeeze-and-excitation block, and a zero-initialized projection that adds the resulting residual to a backbone feature map so that the model behaves as identity at init.

import torch
import torch.nn as nn
import torch.nn.functional as F

class ZeroConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, 
                             stride=stride, padding=padding)
        # Zero initialization of weights and bias
        nn.init.zeros_(self.conv.weight)
        nn.init.zeros_(self.conv.bias)

    def forward(self, x):
        return self.conv(x)
    
class ControlNet(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels, num_blocks=3):
        super().__init__()

        # Initial feature extraction from control signal
        self.init_conv = nn.Conv2d(in_channels, mid_channels, 
                                  kernel_size=3, padding=1)
        
        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.BatchNorm2d(mid_channels),
                nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.BatchNorm2d(mid_channels)
            ) for _ in range(num_blocks)
        ])

        # Residual skip connections
        self.skip_connections = nn.ModuleList([
            nn.Conv2d(mid_channels, mid_channels, kernel_size=1)
            for _ in range(num_blocks)
        ])

        # Channel attention mechanism
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(mid_channels, mid_channels // 4, 1),
            nn.ReLU(),
            nn.Conv2d(mid_channels // 4, mid_channels, 1),
            nn.Sigmoid()
        )

        # Zero-initialized
        self.output_proj = ZeroConv2d(mid_channels, out_channels, kernel_size=1)

    def forward(self, x, control_signal):
        # initial features from control signal
        features = self.init_conv(control_signal)

        # blocks with residual connections
        for block, skip in zip(self.blocks, self.skip_connections):
            residual = skip(features)
            features = block(features)
            features = features + residual
            
            # Apply channel attention
            attention = self.attention(features)
            features = features * attention

        control_output = self.output_proj(features)
        return x + control_output
    
def test_controlnet():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Basic functionality and shape
    model = ControlNet(in_channels=3, mid_channels=64, out_channels=3).to(device)
    x = torch.randn(2, 3, 32, 32).to(device)
    control = torch.randn(2, 3, 32, 32).to(device)

    output = model(x, control)
    assert output.shape == x.shape, f"Shape mismatch: {output.shape} vs {x.shape}"
    
    # Zero initialization 
    x = torch.ones(2, 3, 32, 32).to(device)
    control = torch.zeros(2, 3, 32, 32).to(device)
    output = model(x, control)
    assert torch.allclose(output, x, atol=1e-5), "Zero control should not affect input"
    
    print("All tests passed successfully!")

if __name__ == "__main__":
    test_controlnet()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = ControlNet(in_channels=3, mid_channels=64, out_channels=3).to(device)
    
    x = torch.randn(1, 3, 32, 32).to(device)
    control = torch.randn(1, 3, 32, 32).to(device)
    
    # Generate controlled output
    with torch.no_grad():
        output = model(x, control)
    
    print(f"Input shape: {x.shape}")
    print(f"Control signal shape: {control.shape}")
    print(f"Output shape: {output.shape}")