Minimal Implementation: Independent Multi-Control PixelDiT
This document is the minimal code-level specification for reimplementing the user's method. It intentionally focuses on the smallest set of modules needed to support:
- single-control generation:
depth,seg,edge - multi-control generation:
depth_seg,depth_edge,seg_edge,depth_seg_edge - independent modality branches
- layer-wise gated residual fusion
- strict single-condition hard selection
The current repository implementation lives mainly in:
pixdit_core/pixeldit_t2i_control.py
t2i/diffusion/model/control_trainer.py
t2i/diffusion/data/datasets/control_datasets.py
t2i/train_control.py
t2i/diffusion/losses/multi_condition_cycle.py
t2i/diffusion/losses/edge_cycle.py
1. Core Design
The base model is a PixelDiT text-to-image model. The control extension adds local condition residuals into the transformer blocks.
The key innovation is:
Do not concatenate depth/seg/edge into one shared control encoder.
Use one independent branch per condition, then fuse branch residuals layer-wise.
The control order is fixed everywhere:
CONTROL_NAMES = ("depth", "seg", "edge")
control_keep shape:
[B, 3]
control_keep meanings:
depth only: [1, 0, 0]
seg only: [0, 1, 0]
edge only: [0, 0, 1]
depth + seg: [1, 1, 0]
depth + edge: [1, 0, 1]
seg + edge: [0, 1, 1]
depth + seg+edge: [1, 1, 1]
2. Minimal Model Class
Use the original PixelDiT backbone as parent class. Add independent branch modules only when control_mode == "multi".
import torch
import torch.nn as nn
import torch.nn.functional as F
CONTROL_NAMES = ("depth", "seg", "edge")
class MultiControlPixelDiT(BasePixelDiT):
def __init__(
self,
*,
use_depth_condition=True,
control_mode="multi",
control_names=CONTROL_NAMES,
depth_channels=1,
hidden_size=1536,
depth_base_channels=64,
depth_max_channels=512,
num_inject_layers=14,
init_gate_logits=(0.5, 0.0, -0.5),
enable_structure_inject=True,
control_structure_inject=(True, True, False),
alpha_inject=2.0,
freeze_backbone=True,
freeze_control_branches=(),
pretrained_ckpt=None,
skip_pretrained_modules=(),
**backbone_kwargs,
):
super().__init__(**backbone_kwargs)
self.control_mode = control_mode
self.control_names = tuple(control_names)
self.num_controls = len(self.control_names)
self.depth_channels = depth_channels
self.num_inject_layers = num_inject_layers
self.enable_structure_inject = enable_structure_inject
self.control_structure_inject = tuple(control_structure_inject)
self.alpha_inject = float(alpha_inject)
if control_mode == "single":
# Legacy single-control path. The module names are kept as depth_*
# for compatibility with old depth-control checkpoints, even when
# the input condition is seg or edge in single-control baselines.
self.depth_encoder = DepthEncoder(depth_channels, depth_base_channels, depth_max_channels)
self.depth_adapters = nn.ModuleList([
StructureAwareGatedZeroAdapter(hidden_size)
for _ in range(num_inject_layers)
])
self.seg_encoder = None
self.seg_adapters = None
self.edge_encoder = None
self.edge_adapters = None
self.control_gate_logits = None
elif control_mode == "multi":
# Three independent branches.
self.depth_encoder = DepthEncoder(depth_channels, depth_base_channels, depth_max_channels)
self.depth_adapters = nn.ModuleList([
StructureAwareGatedZeroAdapter(hidden_size)
for _ in range(num_inject_layers)
])
self.seg_encoder = DepthEncoder(depth_channels, depth_base_channels, depth_max_channels)
self.seg_adapters = nn.ModuleList([
StructureAwareGatedZeroAdapter(hidden_size)
for _ in range(num_inject_layers)
])
self.edge_encoder = DepthEncoder(depth_channels, depth_base_channels, depth_max_channels)
self.edge_adapters = nn.ModuleList([
StructureAwareGatedZeroAdapter(hidden_size)
for _ in range(num_inject_layers)
])
# Layer-wise scalar logits: [L, 3].
row = torch.tensor(init_gate_logits, dtype=torch.float32)
self.control_gate_logits = nn.Parameter(
row.view(1, 3).repeat(num_inject_layers, 1).clone()
)
else:
raise ValueError(f"unsupported control_mode={control_mode}")
self.last_gate_weights = None
if pretrained_ckpt is not None:
self.load_pretrained_backbone(pretrained_ckpt, skip_modules=skip_pretrained_modules)
if freeze_backbone:
self.freeze_base_backbone_except_control()
for branch in freeze_control_branches:
self.freeze_control_branch(branch)
3. Structure-Aware Adapter
The adapter must expose compute_residual() so the multi-branch model can compute residuals before adding them to the hidden state.
class StructureAwareGatedZeroAdapter(nn.Module):
def __init__(self, hidden_size):
super().__init__()
self.norm = nn.LayerNorm(hidden_size)
self.proj = nn.Linear(hidden_size, hidden_size)
self.gate = nn.Parameter(torch.zeros(()))
def compute_residual(
self,
cond_tokens: torch.Tensor,
structure_map: torch.Tensor | None = None,
alpha_inject: float = 0.5,
) -> torch.Tensor:
residual = self.gate * self.proj(self.norm(cond_tokens))
if structure_map is not None and alpha_inject != 0.0:
residual = residual * (1.0 + float(alpha_inject) * structure_map.to(residual.dtype))
return residual
def forward(self, x, cond_tokens, structure_map=None, alpha_inject=0.5):
return x + self.compute_residual(cond_tokens, structure_map, alpha_inject)
4. Control Splitting
For multi-control, the model receives control with shape [B, 3, H, W] or [B, 3*C, H, W]. For this project each control is one channel.
def split_controls(control: torch.Tensor, control_keep: torch.Tensor | None):
# control: [B, 3, H, W]
assert control.ndim == 4
b, c, h, w = control.shape
assert c % 3 == 0
ch = c // 3
parts = [
control[:, 0 * ch:1 * ch],
control[:, 1 * ch:2 * ch],
control[:, 2 * ch:3 * ch],
]
if control_keep is None:
control_keep = control.new_ones((b, 3))
return parts, control_keep.to(control.dtype)
5. Gate Weight Function
This is the central single-vs-multi behavior.
Rules:
- If exactly one control is active for a sample, hard select it.
- If more than one control is active, softmax only over active controls.
- Inactive controls receive weight zero.
- Gate logits are used only for multi-condition samples.
def per_layer_per_sample_weights(
gate_logits: torch.Tensor, # [L, 3]
keep: torch.Tensor, # [B, 3]
) -> torch.Tensor:
l, n = gate_logits.shape
b = keep.shape[0]
assert n == keep.shape[1] == 3
active_count = keep.sum(dim=1) # [B]
weights = gate_logits.new_zeros((l, b, n))
# Single-condition: hard select. Gate is ignored.
single = active_count == 1
if single.any():
weights[:, single, :] = keep[single].view(1, -1, n)
# Multi-condition: masked softmax over active controls.
multi = active_count > 1
if multi.any():
logits = gate_logits[:, None, :].expand(l, int(multi.sum()), n)
mask = keep[multi].bool().view(1, -1, n)
logits = logits.masked_fill(~mask, -torch.finfo(logits.dtype).max)
weights[:, multi, :] = torch.softmax(logits, dim=-1)
return weights # [L, B, 3]
This produces the fusion formula:
R_l = w_l,b,depth * R_l,depth
+ w_l,b,seg * R_l,seg
+ w_l,b,edge * R_l,edge
6. Branch Feature Computation
Each branch is computed independently. For efficiency and strict gradient behavior, a branch can be skipped if no sample in the batch uses it.
def compute_branch_features(self, controls, keep, grid_hw):
encoders = [self.depth_encoder, self.seg_encoder, self.edge_encoder]
adapters = [self.depth_adapters, self.seg_adapters, self.edge_adapters]
branch_tokens = [None, None, None]
branch_structs = [None, None, None]
for i, control_i in enumerate(controls):
if keep[:, i].sum() <= 0:
continue
# Encoder returns one token tensor per injection layer.
feats_i = encoders[i](control_i) # list length L, each [B, T, D]
branch_tokens[i] = feats_i
use_structure = (
self.enable_structure_inject
and self.alpha_inject != 0.0
and self.control_structure_inject[i]
)
if use_structure:
branch_structs[i] = [sobel_structure_map(control_i, grid_hw) for _ in range(self.num_inject_layers)]
return branch_tokens, branch_structs
7. Forward Fusion Pseudocode
Inside the PixelDiT block loop, inject the fused residual at each target layer.
def forward(self, x, t, y, *, control=None, control_keep=None, **kwargs):
if self.control_mode == "single":
# Legacy path: one control branch only.
cond_feats = self.depth_encoder(control)
for layer_idx, block in enumerate(self.blocks):
x = block(x, t, y)
if layer_idx in self.inject_layer_indices:
j = self.inject_layer_indices.index(layer_idx)
struct = sobel_structure_map(control, grid_hw) if self.enable_structure_inject else None
x = self.depth_adapters[j](x, cond_feats[j], struct, self.alpha_inject)
return x
# Multi-control path.
controls, keep = split_controls(control, control_keep)
branch_feats, branch_structs = self.compute_branch_features(controls, keep, grid_hw)
weights = per_layer_per_sample_weights(self.control_gate_logits.to(x.dtype), keep)
self.last_gate_weights = weights.detach().float().cpu()
for layer_idx, block in enumerate(self.blocks):
x = block(x, t, y)
if layer_idx not in self.inject_layer_indices:
continue
j = self.inject_layer_indices.index(layer_idx)
fused = 0.0
for branch_idx in range(3):
if branch_feats[branch_idx] is None:
continue
adapter = [self.depth_adapters, self.seg_adapters, self.edge_adapters][branch_idx][j]
cond_j = branch_feats[branch_idx][j]
struct_j = None if branch_structs[branch_idx] is None else branch_structs[branch_idx][j]
residual = adapter.compute_residual(cond_j, struct_j, self.alpha_inject)
w = weights[j, :, branch_idx].view(-1, 1, 1)
fused = fused + w * residual
x = x + fused
return x
8. Training Mode Sampling
The training loop samples one mode per step:
CONTROL_MODES = [
"depth",
"seg",
"edge",
"depth_seg",
"depth_edge",
"seg_edge",
"depth_seg_edge",
]
Final mixed probabilities:
CONTROL_PROBS = [0.15, 0.15, 0.15, 0.12, 0.12, 0.12, 0.19]
DDP ranks must use the same mode:
def sample_control_mode(modes, probs, device):
probs = torch.tensor(probs, dtype=torch.float32, device=device)
probs = probs / probs.sum()
idx = torch.multinomial(probs, 1)
if torch.distributed.is_available() and torch.distributed.is_initialized():
torch.distributed.broadcast(idx, src=0)
return modes[int(idx.item())]
Mode to keep mask:
def mode_to_keep(mode: str):
tokens = set(mode.split("_"))
return torch.tensor([
1.0 if "depth" in tokens else 0.0,
1.0 if "seg" in tokens else 0.0,
1.0 if "edge" in tokens else 0.0,
])
Apply mode:
def apply_multi_control_mode(control, mode):
# control: [B, 3, H, W]
keep = mode_to_keep(mode).to(control.device, control.dtype)
keep_b = keep.view(1, 3, 1, 1)
control = control * keep_b
control_keep = keep.view(1, 3).expand(control.shape[0], 3)
return control, control_keep
9. Gradient Masking Invariant
For single active modes, only the active branch should update. For multi-condition, only active branches plus the gate update.
def mask_inactive_control_grads(model, control_mode: str):
tokens = set(control_mode.split("_"))
active = {
"depth": "depth" in tokens,
"seg": "seg" in tokens,
"edge": "edge" in tokens,
}
gate_active = sum(active.values()) > 1
for name, p in model.named_parameters():
if p.grad is None:
continue
if "control_gate_logits" in name and not gate_active:
p.grad = None
elif ("depth_encoder" in name or "depth_adapters" in name) and not active["depth"]:
p.grad = None
elif ("seg_encoder" in name or "seg_adapters" in name) and not active["seg"]:
p.grad = None
elif ("edge_encoder" in name or "edge_adapters" in name) and not active["edge"]:
p.grad = None
Important exception: for the legacy single-control baseline path, parameters are still named depth_encoder/depth_adapters even if the condition is seg or edge. Do not apply this independent-branch gradient mask to the single-control baseline model.
10. Dataset Output Contract
For three-control training, each item should include:
data_info = {
"control": torch.cat([depth, seg, edge], dim=0), # [3, H, W]
"control_keep": torch.tensor([1.0, 1.0, 1.0]),
"control_mode": "depth_seg_edge",
"depth": depth,
"seg": seg,
"edge": edge,
}
The training loop samples the active mode and zeroes inactive channels.
For single-control baselines, each item should include:
data_info = {
"control": control, # [1, H, W]
"control_keep": torch.tensor([1.0]),
"control_mode": "depth" or "seg" or "edge",
}
11. Cycle Loss Minimal Interface
The multi-condition cycle wrapper receives generated image and condition labels.
class MultiConditionCycleLoss(nn.Module):
def __init__(self, depth_cycle_loss=None, seg_cycle_loss=None, edge_cycle_loss=None,
depth_weight=1.0, seg_weight=1.0, edge_weight=1.0):
super().__init__()
self.depth_cycle_loss = depth_cycle_loss
self.seg_cycle_loss = seg_cycle_loss
self.edge_cycle_loss = edge_cycle_loss
self.depth_weight = depth_weight
self.seg_weight = seg_weight
self.edge_weight = edge_weight
def forward(self, gen_image, depth_01=None, seg_01=None, gt_image_m11=None, control_mode="depth_seg"):
tokens = set(control_mode.split("_"))
total = gen_image.new_zeros(())
if "depth" in tokens and self.depth_cycle_loss is not None:
total = total + self.depth_weight * self.depth_cycle_loss(gen_image, depth_01)
if "seg" in tokens and self.seg_cycle_loss is not None:
total = total + self.seg_weight * self.seg_cycle_loss(gen_image, seg_01)
if "edge" in tokens and self.edge_cycle_loss is not None:
total = total + self.edge_weight * self.edge_cycle_loss(gen_image, gt_image_m11)
return total
Depth and seg compare generated-image-derived structure to the condition label. Edge uses generated RGB and GT RGB because offline Canny labels are threshold-sensitive.
12. SoftCanny Edge Cycle
Minimal idea:
class SoftCannyImagePyramidCycleLoss(nn.Module):
def forward(self, gen_image_m11, gt_image_m11):
threshold = uniform(0.2745, 0.5882)
gen_edge = soft_canny(gen_image_m11, threshold)
gt_edge = soft_canny(gt_image_m11, threshold).detach()
return smooth_l1(gen_edge, gt_edge)
Current parameters:
gaussian_kernel: 11
threshold_min: 0.2745
threshold_max: 0.5882
temperature: 0.03
cycle_scales: [512, 256, 128, 64]
cycle_scale_weights: [0.1, 0.25, 1.0, 0.25]
13. Minimal Final Config Values
Final mixed-control config:
control_names: [depth, seg, edge]
init_gate_logits: [0.5, 0.0, -0.5]
control_structure_inject: [true, true, false]
freeze_backbone: true
pretrained_ckpt: /media/home/songmeixi_insta360.com/PixelDiT-master/t2i/universal_pix_t2i_workdirs/exp_pixeldit_threecontrol_v1_mixed_smalllr_from_softcanny2k/checkpoints/epoch_1_step_2000.pth
depth_branch_lr_scale: 0.05
seg_branch_lr_scale: 0.1
edge_branch_lr_scale: 0.1
gate_lr_scale: 0.5
cycle_weight: 0.005
control_probs: [0.15, 0.15, 0.15, 0.12, 0.12, 0.12, 0.19]
14. Critical Implementation Checks
Before considering an implementation correct:
1. depth-only output ignores gate and uses only depth branch.
2. seg-only output ignores gate and uses only seg branch.
3. edge-only output ignores gate and uses only edge branch.
4. multi-condition output uses masked softmax over active controls only.
5. inactive branches receive no gradients.
6. gate receives gradients only for multi-condition samples.
7. DDP ranks sample the same control_mode each step.
8. edge structure injection is disabled in final mixed training.
9. corrupt edge maps are skipped/replaced, not silently loaded.
10. gate weights are logged for interpretability.