feat(lora): parallel decoding heads for pdd acceleration files

PDD files pair a backbone LoRA with the output projections repeated per
interval of a training grid; each step fuses the heads of its block
into one projection. network_pdd reads the grid from the metadata,
swaps a ParallelHead in for each projection, fuses from the scheduler's
step_index and pins the step count and shift while heads are
installed. The MiniMax loader declares which scheduler each head
follows.

modules/lora/native_adapter.py

@@ -1140,6 +1140,7 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders):
             net = sub
         else:
             net.modules.update(sub.modules)
+            net.extras.update(sub.extras)
     sd_models_utils.state_dict_cache.disable()
     if net is not None and mismatch > 0: # applying only the layers that fit leaves the model in a state nothing was trained for
         log.error(f'Network load: type=LoRA name="{name}" modules={len(net.modules)} mismatch={mismatch} shapes do not match the loaded model')

modules/lora/network.py

@@ -161,6 +161,7 @@ class Network:  # LoraModule
         self.block_spec = None # raw lbw= value; per-layer factors resolve through lora_blocks
         self.pending_config = None # staged multipliers; network_activate promotes them after the removal pass so fuse removal subtracts the delta that was applied
         self.modules = {}
+        self.extras = {} # non-delta payloads a family carries, e.g. parallel heads
         self.mismatch = 0 # deltas dropped for not fitting their target module; try_load_chain refuses the file when non-zero
         self.bundle_embeddings = {}
         self.mtime = None

modules/lora/network_pdd.py

@@ -0,0 +1,305 @@
+"""Parallel decoding distillation heads carried by a native network.
+
+A PDD file pairs a backbone LoRA with the output projections repeated once per interval of an N-step
+training grid; each step fuses the heads of its block into one projection, so N / block_size evaluations
+walk the trajectory. Heads ride on ``Network.extras['pdd']``: ``reconcile`` installs and removes them with
+the loaded set, ``pin`` holds the step count and schedule the file was distilled for.
+"""
+
+import copy
+import weakref
+import torch
+from modules.logger import log
+
+
+METADATA_STEPS = 'pdd_num_steps'
+METADATA_BLOCK = 'pdd_block_size'
+EXTRAS_KEY = 'pdd'
+
+
+class ArchSpec:
+    """How an architecture hosts parallel heads: the scheduler behind each head and how interval counts map onto its num_inference_steps."""
+
+    def __init__(self, schedulers=None, default_scheduler='scheduler', steps_for=None):
+        self.schedulers = schedulers or {} # head path -> attribute of the scheduler the head was trained on
+        self.default_scheduler = default_scheduler
+        self.steps_for = steps_for or (lambda intervals: intervals) # num_inference_steps that yields this many grid intervals
+
+    def scheduler_name(self, head):
+        return self.schedulers.get(head, self.default_scheduler)
+
+
+class ParallelHeads:
+    """The head tensors of one file and the grid they were trained on."""
+
+    def __init__(self, num_steps, block_size, heads):
+        self.num_steps = num_steps
+        self.block_size = block_size
+        self.heads = heads # head path -> (weight [N, out, in], bias [N, out] or None)
+        self.nfe = num_steps // block_size
+
+
+class Installed:
+    """Bookkeeping for the heads currently swapped into a pipeline."""
+
+    def __init__(self, name, strength, component, modules, heads, spec):
+        self.name = name
+        self.strength = strength
+        self.component = component
+        self.modules = modules # head path -> (parent, attribute, original module)
+        self.heads = heads
+        self.spec = spec
+        self.steps = spec.steps_for(heads.nfe)
+
+
+def detect(metadata):
+    """The (num_steps, block_size) grid a file declares; None without PDD metadata, ValueError for an unusable grid."""
+    metadata = metadata or {}
+    if METADATA_STEPS not in metadata:
+        return None
+    num_steps = int(metadata[METADATA_STEPS])
+    block_size = int(metadata.get(METADATA_BLOCK, 1))
+    if num_steps < 1 or block_size < 1 or num_steps % block_size != 0:
+        raise ValueError(f'grid={num_steps} block={block_size}')
+    return num_steps, block_size
+
+
+def load(name, metadata, state_dict):
+    """Collect the per-interval head tensors of a file, or None when the file carries no PDD grid."""
+    try:
+        grid = detect(metadata)
+    except ValueError as e:
+        log.error(f'Network load: type=PDD name="{name}" {e} block size must divide the grid')
+        return None
+    if grid is None:
+        return None
+    num_steps, block_size = grid
+    heads = {}
+    for key, tensor in state_dict.items():
+        if key.endswith('.weight') and tensor.ndim == 3 and tensor.shape[0] == num_steps:
+            path = key[:-len('.weight')]
+            heads[path] = (tensor, state_dict.get(f'{path}.bias', None))
+    if len(heads) == 0:
+        log.error(f'Network load: type=PDD name="{name}" grid={num_steps} block={block_size} no head tensors')
+        return None
+    log.debug(f'Network load: type=PDD name="{name}" grid={num_steps} block={block_size} nfe={num_steps // block_size} heads={list(heads)}')
+    return ParallelHeads(num_steps, block_size, heads)
+
+
+def try_load(name, network_on_disk, lora_scale): # pylint: disable=unused-argument
+    """Family loader for the native chain: a network carrying only the heads."""
+    metadata = getattr(network_on_disk, 'metadata', None) or {}
+    if METADATA_STEPS not in metadata:
+        return None
+    from modules.lora import native_adapter
+    state_dict = native_adapter.read_state_dict(network_on_disk.filename, what='network')
+    heads = load(name, metadata, state_dict)
+    if heads is None:
+        return None
+    net = native_adapter.new_network(name, network_on_disk)
+    net.extras[EXTRAS_KEY] = heads
+    return net
+
+
+def base_tensors(module):
+    """Float copies of a projection's weight and bias for the strength blend; None when the weight is not a plain tensor."""
+    weight = getattr(module, 'weight', None)
+    if weight is None or getattr(module, 'sdnq_dequantizer', None) is not None or not torch.is_floating_point(weight) or weight.ndim != 2:
+        return None, None
+    bias = getattr(module, 'bias', None)
+    return weight.detach().to(dtype=torch.float32), None if bias is None else bias.detach().to(dtype=torch.float32)
+
+
+class ParallelHead(torch.nn.Module):
+    """An output projection replaced by its per-interval heads, fused per step for the block the scheduler is about to take."""
+
+    def __init__(self, base, weight, bias, strength, get_scheduler, block_size, intervals):
+        super().__init__()
+        object.__setattr__(self, 'base', base) # kept out of the module tree so nothing walks, offloads or serializes it
+        object.__setattr__(self, 'get_scheduler', get_scheduler)
+        self.weight = torch.nn.Parameter(weight.to(dtype=torch.float32), requires_grad=False) # float32 like the projection it replaces
+        self.bias = None if bias is None else torch.nn.Parameter(bias.to(dtype=torch.float32), requires_grad=False)
+        self.register_buffer('intervals', intervals.to(dtype=torch.float32))
+        self.in_features = weight.shape[2]
+        self.out_features = weight.shape[1]
+        self.num_steps = weight.shape[0]
+        self.block_size = block_size
+        self.strength = strength
+        self.base_weight, self.base_bias = base_tensors(base) if strength != 1.0 else (None, None)
+        self.fused_index = None
+        self.fused_weight = None
+        self.fused_bias = None
+        self.overflow_warned = False
+
+    def step_index(self):
+        scheduler = self.get_scheduler()
+        index = getattr(scheduler, 'step_index', None) if scheduler is not None else None
+        index = 0 if index is None else int(index)
+        nfe = self.num_steps // self.block_size
+        if index >= nfe:
+            if not self.overflow_warned:
+                self.overflow_warned = True
+                log.warning(f'Network: type=PDD step={index} nfe={nfe} schedule longer than the distilled grid')
+            index = nfe - 1
+        return index
+
+    def fuse(self, index):
+        start = index * self.block_size
+        stop = start + self.block_size
+        plan = torch.zeros(self.num_steps, dtype=torch.float32, device=self.weight.device)
+        span = self.intervals[start:stop]
+        plan[start:stop] = (span / span.sum()).to(device=plan.device)
+        weight = torch.tensordot(plan, self.weight.detach(), dims=1)
+        bias = None if self.bias is None else plan @ self.bias.detach()
+        if self.base_weight is not None:
+            base_weight = self.base_weight.to(device=weight.device)
+            weight = base_weight + self.strength * (weight - base_weight)
+            if bias is not None and self.base_bias is not None:
+                base_bias = self.base_bias.to(device=bias.device)
+                bias = base_bias + self.strength * (bias - base_bias)
+        self.fused_index, self.fused_weight, self.fused_bias = index, weight, bias
+
+    def forward(self, hidden_states):
+        index = self.step_index()
+        if index != self.fused_index or self.fused_weight is None or self.fused_weight.device != self.weight.device:
+            self.fuse(index)
+        weight = self.fused_weight.to(device=hidden_states.device, dtype=hidden_states.dtype)
+        bias = None if self.fused_bias is None else self.fused_bias.to(device=hidden_states.device, dtype=hidden_states.dtype)
+        return torch.nn.functional.linear(hidden_states, weight, bias)
+
+
+def grid_intervals(scheduler, num_steps, spec):
+    """Interval lengths of the training grid in ascending time, from a pristine scheduler copy immune to a live shift override."""
+    probe = scheduler.__class__.from_config(scheduler.config) if hasattr(scheduler, 'from_config') else copy.deepcopy(scheduler)
+    probe.set_timesteps(spec.steps_for(num_steps))
+    sigmas = probe.sigmas.detach().to(device='cpu', dtype=torch.float64)
+    if sigmas.numel() != num_steps + 1:
+        return None
+    return (1.0 - sigmas).diff()
+
+
+def submodule(component, path):
+    try:
+        return component.get_submodule(path)
+    except AttributeError:
+        return None
+
+
+def owner(pipe, heads, components):
+    """The component holding every head projection, as (name, module); (None, None) when no component has them all."""
+    for name in components:
+        component = getattr(pipe, name, None)
+        if component is not None and hasattr(component, 'get_submodule') and all(submodule(component, path) is not None for path in heads.heads):
+            return name, component
+    return None, None
+
+
+def target_shape(module):
+    dequantizer = getattr(module, 'sdnq_dequantizer', None)
+    if dequantizer is not None and getattr(dequantizer, 'original_shape', None) is not None:
+        return tuple(dequantizer.original_shape)
+    weight = getattr(module, 'weight', None)
+    return tuple(weight.shape) if weight is not None else None
+
+
+def install(pipe, net, heads, spec, components):
+    """Swap the heads into the component that owns their projections; True when the module tree changed."""
+    component_name, component = owner(pipe, heads, components)
+    if component is None:
+        log.error(f'Network load: type=PDD name="{net.name}" heads={list(heads.heads)} no loaded component holds these projections')
+        return False
+    strength = float(net.te_multiplier) # transformer-keyed layers scale by the te multiplier, see network.NetworkModule.multiplier
+    pipe_ref = weakref.ref(pipe)
+    modules = {}
+    for path, (weight, bias) in heads.heads.items():
+        module = submodule(component, path)
+        shape = target_shape(module)
+        if shape != tuple(weight.shape[1:]):
+            log.error(f'Network load: type=PDD name="{net.name}" head={path} shape={list(weight.shape[1:])} module={list(shape) if shape else None} shape mismatch')
+            for parent, attr, original in modules.values():
+                setattr(parent, attr, original)
+            return False
+        scheduler_name = spec.scheduler_name(path)
+        scheduler = getattr(pipe, scheduler_name, None)
+        intervals = grid_intervals(scheduler, heads.num_steps, spec) if scheduler is not None else None
+        if intervals is None:
+            log.error(f'Network load: type=PDD name="{net.name}" head={path} scheduler={scheduler.__class__.__name__} cannot build a {heads.num_steps}-interval grid')
+            for parent, attr, original in modules.values():
+                setattr(parent, attr, original)
+            return False
+        def get_scheduler(name=scheduler_name):
+            owner_pipe = pipe_ref()
+            return getattr(owner_pipe, name, None) if owner_pipe is not None else None
+        head = ParallelHead(module, weight, bias, strength, get_scheduler, heads.block_size, intervals)
+        parent_path, _, attr = path.rpartition('.')
+        parent = component.get_submodule(parent_path) if parent_path else component
+        setattr(parent, attr, head)
+        modules[path] = (parent, attr, module)
+    pipe.sdnext_pdd = Installed(net.name, strength, component_name, modules, heads, spec)
+    log.info(f'Network load: type=PDD name="{net.name}" component={component_name} heads={list(modules)} grid={heads.num_steps} block={heads.block_size} nfe={heads.nfe} steps={pipe.sdnext_pdd.steps} strength={strength}')
+    return True
+
+
+def restore(pipe):
+    """Put the original projections back; True when heads were installed."""
+    state = getattr(pipe, 'sdnext_pdd', None)
+    if state is None:
+        return False
+    for parent, attr, original in state.modules.values():
+        setattr(parent, attr, original)
+    del pipe.sdnext_pdd
+    log.info(f'Network unload: type=PDD name="{state.name}" component={state.component} heads={list(state.modules)}')
+    return True
+
+
+def arch_spec():
+    """The parallel-head spec of the loaded architecture's native loader module, or None."""
+    import importlib
+    from modules import shared
+    from modules.lora import lora_load
+    module_name = lora_load.NATIVE_DISPATCH.get(shared.sd_model_type)
+    if module_name is None:
+        return None
+    return getattr(importlib.import_module(module_name), 'PDD', None)
+
+
+def reconcile(pipe, loaded, components):
+    """Match the installed heads to the loaded networks; True when the module tree changed."""
+    carriers = [net for net in loaded if EXTRAS_KEY in getattr(net, 'extras', {})]
+    if len(carriers) == 0:
+        return restore(pipe)
+    if len(carriers) > 1:
+        log.warning(f'Network load: type=PDD networks={[net.name for net in carriers]} one grid per model, using first')
+    net = carriers[0]
+    current = getattr(pipe, 'sdnext_pdd', None)
+    if current is not None and current.name == net.name and current.strength == float(net.te_multiplier):
+        return False
+    spec = arch_spec()
+    if spec is None:
+        from modules import shared
+        log.error(f'Network load: type=PDD name="{net.name}" type={shared.sd_model_type} architecture has no parallel head support')
+        return restore(pipe)
+    changed = restore(pipe)
+    return install(pipe, net, net.extras[EXTRAS_KEY], spec, components) or changed
+
+
+def pin(p, model):
+    """Hold a generation on the distilled step count and shipped schedule while heads are installed; returns the pinned step count or None."""
+    state = getattr(model, 'sdnext_pdd', None)
+    if state is None:
+        return None
+    shifts = {}
+    for path in state.modules:
+        name = state.spec.scheduler_name(path)
+        scheduler = getattr(model, name, None)
+        if scheduler is not None and hasattr(scheduler, 'set_shift') and getattr(scheduler, 'config', None) is not None and 'shift' in scheduler.config:
+            scheduler.set_shift(scheduler.config['shift'])
+            shifts[name] = scheduler.config['shift']
+    requested = p.steps
+    p.steps = state.steps
+    if getattr(p, 'task_args', None) is not None:
+        p.task_args['num_inference_steps'] = state.steps
+    if getattr(model, 'num_timesteps', None) is not None:
+        model.num_timesteps = state.heads.nfe # the progress total counts transformer evaluations
+    log.info(f'Network: type=PDD name="{state.name}" steps={state.steps} requested={requested} nfe={state.heads.nfe} shift={shifts}')
+    return state.steps

modules/lora/networks.py

@@ -48,8 +48,9 @@ from modules.lora import lora_common as l
 from modules.lora import lora_overrides
 from modules.lora import lora_sdnq
 from modules.lora import lora_stack
+from modules.lora import network_pdd
 from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights
-from modules import shared, devices, sd_models
+from modules import shared, devices, sd_models, errors
 from modules.logger import log, console
 
 
@@ -334,7 +335,12 @@ def finish_pass(ctx, t0):
         log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={ctx.applied_weight} bias={ctx.applied_bias} refused={ctx.refused} network partially applied')
     if l.debug and len(l.loaded_networks) > 0:
         log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={ctx.active_components} layers={ctx.total} weights={ctx.applied_weight} bias={ctx.applied_bias} refused={ctx.refused} backup={round(ctx.backup_size/1024/1024/1024, 2)} fuse={ctx.fuse}:{shared.opts.lora_fuse_diffusers} device={ctx.device} time={l.timer.summary}')
-    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(ctx.group_stripped) > 0:
+    try:
+        heads_changed = network_pdd.reconcile(ctx.sd_model, l.loaded_networks, default_components) # carried projections swap with the loaded set before the offload snapshot below
+    except Exception as e:
+        heads_changed = False
+        errors.display(e, 'Network load: type=PDD')
+    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(ctx.group_stripped) > 0 or heads_changed:
         sd_models.set_diffuser_offload(ctx.sd_model, op="model")
 
 

modules/processing_args.py

@@ -12,6 +12,7 @@ from modules.attention import context as attention_context
 from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p
 from modules.processing_helpers import get_generator, apply_circular # pylint: disable=unused-import
 from modules.processing_prompt import set_prompt
+from modules.lora import network_pdd
 from modules.api import helpers
 
 
@@ -263,6 +264,10 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
 
     possible = get_params(model)
 
+    pinned = network_pdd.pin(p, model) # installed parallel-decoding heads fix the step count and schedule
+    if pinned is not None and 'num_inference_steps' in possible:
+        kwargs['num_inference_steps'] = pinned
+
     debug_log(f'Pipeline: cls={cls} possible={possible}')
     steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1']))
     clip_skip = kwargs.pop("clip_skip", 1)

pipelines/minimax/minimax_lora.py

@@ -15,7 +15,11 @@ import re
 import torch
 
 from modules.logger import log
-from modules.lora import native_adapter
+from modules.lora import native_adapter, network_pdd
+
+
+# Parallel decoding heads: the audio projection follows the audio schedule, and MiniMaxH3Scheduler counts the terminal sigma in num_inference_steps.
+PDD = network_pdd.ArchSpec(schedulers={"audio_proj_out": "audio_scheduler"}, steps_for=lambda intervals: intervals + 1)
 
 
 KNOWN_PREFIXES = (
@@ -199,8 +203,9 @@ def network_prefix_for(prefix_used):
 
 
 def file_alpha(network_on_disk):
-    """The training alpha some trainers record in the safetensors metadata instead of per-key tensors, or None."""
-    alpha = (getattr(network_on_disk, "metadata", None) or {}).get("alpha")
+    """The file-level training alpha from the safetensors metadata (alpha, or lora_alpha in PDD files), or None."""
+    metadata = getattr(network_on_disk, "metadata", None) or {}
+    alpha = metadata.get("alpha", metadata.get("lora_alpha"))
     if alpha is None:
         return None
     try:
@@ -287,5 +292,6 @@ def try_load(name, network_on_disk, lora_scale):
         family_loaders=(
             try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
             try_load_ia3, try_load_glora, try_load_norm, try_load_full,
+            network_pdd.try_load,
         ),
     )

test/test-pdd.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for modules.lora.network_pdd.
+
+Checks the parallel decoding head mechanics against the reference formulas
+shipped with alibaba-pai/MiniMax-H3-Acc-LoRAs (minimax_h3_pdd.py):
+
+- ``detect`` / ``load`` metadata and head-tensor discovery
+- ``grid_intervals`` against ``pdd_time_grid`` for the video and audio shifts
+- ``ParallelHead`` against ``MiniMaxH3ParallelHead`` on every block, plus the strength blend
+- ``install`` / ``restore`` / ``reconcile`` round trips on a stub pipeline
+- ``pin`` step and shift override
+
+No running server required.
+
+Usage:
+    python test/test-pdd.py
+"""
+
+import os
+import sys
+import types
+
+import torch
+
+script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, script_dir)
+os.chdir(script_dir)
+
+os.environ['SD_INSTALL_QUIET'] = '1'
+
+# Bootstrap cmd_args before any module that pulls in shared.py.
+import modules.cmd_args  # pylint: disable=wrong-import-position
+import installer  # pylint: disable=wrong-import-position
+_orig_argv = sys.argv
+sys.argv = [sys.argv[0]]
+try:
+    modules.cmd_args.parse_args()
+finally:
+    sys.argv = _orig_argv
+installer.add_args(modules.cmd_args.parser)
+modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
+
+import diffusers  # pylint: disable=wrong-import-position
+from modules import shared  # pylint: disable=wrong-import-position,unused-import # shared must initialize before sd_models, which imports back into it
+from modules.errors import log  # pylint: disable=wrong-import-position
+from modules.lora import network_pdd  # pylint: disable=wrong-import-position
+from pipelines.minimax import minimax_lora  # pylint: disable=wrong-import-position
+
+
+NUM_STEPS = 32
+BLOCK = 4
+HIDDEN = 16
+VIDEO_OUT = 8
+AUDIO_OUT = 6
+METADATA = {'pdd_num_steps': '32', 'pdd_block_size': '4', 'lora_rank': '64', 'lora_alpha': '64.0'}
+
+
+# ============================================================
+# Reference implementation (minimax_h3_pdd.py)
+# ============================================================
+
+def reference_time_grid(shift, num_steps):
+    sigma = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)
+    return 1.0 - shift * sigma / (1 + (shift - 1) * sigma)
+
+
+def reference_plan(step_sizes, start, block_size):
+    plan = torch.zeros(1, step_sizes.shape[0], dtype=step_sizes.dtype)
+    span = step_sizes[start:start + block_size].sum()
+    plan[0, start:start + block_size] = step_sizes[start:start + block_size] / span
+    return plan
+
+
+class ReferenceHead(torch.nn.Module):
+    def __init__(self, weight, bias):
+        super().__init__()
+        self.num_steps = weight.shape[0]
+        self.weight = torch.nn.Parameter(weight.clone())
+        self.bias = torch.nn.Parameter(bias.clone())
+        self.plan = torch.zeros(1, self.num_steps)
+
+    def forward(self, hidden_states):
+        plan = self.plan.to(device=self.weight.device, dtype=self.weight.dtype)
+        weight = torch.einsum('pn,noi->poi', plan, self.weight).flatten(0, 1)
+        bias = torch.einsum('pn,no->po', plan, self.bias).flatten()
+        return torch.nn.functional.linear(hidden_states, weight, bias)
+
+
+# ============================================================
+# Test infrastructure
+# ============================================================
+
+results: dict[str, dict] = {}
+
+
+def category(name: str):
+    if name not in results:
+        results[name] = {'passed': 0, 'failed': 0, 'tests': []}
+    return name
+
+
+def record(cat: str, passed: bool, name: str, detail: str = ''):
+    status = 'PASS' if passed else 'FAIL'
+    results[cat]['passed' if passed else 'failed'] += 1
+    results[cat]['tests'].append((status, name))
+    msg = f'  {status}: {name}'
+    if detail:
+        msg += f' ({detail})'
+    if passed:
+        log.info(msg)
+    else:
+        log.error(msg)
+
+
+def run_test(cat: str, fn):
+    name = fn.__name__
+    try:
+        ok = fn()
+        record(cat, ok is not False, name)
+    except AssertionError as e:
+        record(cat, False, name, str(e))
+    except Exception as e:  # pylint: disable=broad-except
+        record(cat, False, name, f'{type(e).__name__}: {e}')
+
+
+# ============================================================
+# Fixtures
+# ============================================================
+
+class StubPipe:
+    """The parts of a modular pipeline the engine touches: components and two schedulers."""
+
+    def __init__(self):
+        self.transformer = torch.nn.Module()
+        self.transformer.proj_out = torch.nn.Linear(HIDDEN, VIDEO_OUT)
+        self.transformer.audio_proj_out = torch.nn.Linear(HIDDEN, AUDIO_OUT)
+        self.scheduler = diffusers.MiniMaxH3Scheduler(shift=12.0)
+        self.audio_scheduler = diffusers.MiniMaxH3Scheduler(shift=3.0)
+        self.num_timesteps = 29
+
+
+def make_heads():
+    torch.manual_seed(0)
+    return network_pdd.ParallelHeads(NUM_STEPS, BLOCK, {
+        'proj_out': (torch.randn(NUM_STEPS, VIDEO_OUT, HIDDEN), torch.randn(NUM_STEPS, VIDEO_OUT)),
+        'audio_proj_out': (torch.randn(NUM_STEPS, AUDIO_OUT, HIDDEN), torch.randn(NUM_STEPS, AUDIO_OUT)),
+    })
+
+
+def make_net(heads, strength=1.0):
+    return types.SimpleNamespace(name='pdd-test', te_multiplier=strength, extras={network_pdd.EXTRAS_KEY: heads})
+
+
+# ============================================================
+# Tests: detection and loading
+# ============================================================
+
+def test_detect_grid():
+    assert network_pdd.detect(METADATA) == (32, 4)
+    assert network_pdd.detect({'pdd_num_steps': '16'}) == (16, 1)
+    assert network_pdd.detect({}) is None
+    assert network_pdd.detect(None) is None
+
+
+def test_detect_rejects_bad_block():
+    try:
+        network_pdd.detect({'pdd_num_steps': '32', 'pdd_block_size': '5'})
+    except ValueError:
+        return True
+    raise AssertionError('block size 5 accepted for a 32 grid')
+
+
+def test_load_collects_heads():
+    state_dict = {
+        'proj_out.weight': torch.zeros(NUM_STEPS, VIDEO_OUT, HIDDEN),
+        'proj_out.bias': torch.zeros(NUM_STEPS, VIDEO_OUT),
+        'audio_proj_out.weight': torch.zeros(NUM_STEPS, AUDIO_OUT, HIDDEN),
+        'audio_proj_out.bias': torch.zeros(NUM_STEPS, AUDIO_OUT),
+        'transformer_blocks.0.attn.to_q.lora_down': torch.zeros(64, HIDDEN),
+        'transformer_blocks.0.attn.to_q.lora_up': torch.zeros(HIDDEN, 64),
+    }
+    heads = network_pdd.load('x', METADATA, state_dict)
+    assert heads is not None and set(heads.heads) == {'proj_out', 'audio_proj_out'}, f'heads={None if heads is None else list(heads.heads)}'
+    assert heads.nfe == 8 and heads.block_size == 4
+    assert heads.heads['proj_out'][1] is not None, 'bias not paired'
+
+
+def test_load_without_metadata_is_none():
+    assert network_pdd.load('x', {'alpha': '1'}, {'proj_out.weight': torch.zeros(NUM_STEPS, VIDEO_OUT, HIDDEN)}) is None
+
+
+def test_load_without_heads_is_none():
+    assert network_pdd.load('x', METADATA, {'transformer_blocks.0.attn.to_q.lora_down': torch.zeros(64, HIDDEN)}) is None
+
+
+# ============================================================
+# Tests: grid and fusion math
+# ============================================================
+
+def test_grid_intervals_match_reference():
+    for shift in (12.0, 3.0):
+        scheduler = diffusers.MiniMaxH3Scheduler(shift=shift)
+        scheduler.set_shift(4.0) # a live override must not leak into the training grid
+        intervals = network_pdd.grid_intervals(scheduler, NUM_STEPS, minimax_lora.PDD)
+        reference = reference_time_grid(shift, NUM_STEPS).diff()
+        assert intervals is not None and intervals.shape == reference.shape, f'shift={shift} shape={None if intervals is None else intervals.shape}'
+        assert torch.allclose(intervals, reference, atol=1e-6), f'shift={shift} maxdiff={(intervals - reference).abs().max().item()}'
+
+
+def test_steps_for_counts_terminal_sigma():
+    assert minimax_lora.PDD.steps_for(8) == 9
+    assert minimax_lora.PDD.scheduler_name('audio_proj_out') == 'audio_scheduler'
+    assert minimax_lora.PDD.scheduler_name('proj_out') == 'scheduler'
+
+
+def test_parallel_head_matches_reference_per_block():
+    heads = make_heads()
+    weight, bias = heads.heads['proj_out']
+    base = torch.nn.Linear(HIDDEN, VIDEO_OUT)
+    scheduler = diffusers.MiniMaxH3Scheduler(shift=12.0)
+    intervals = network_pdd.grid_intervals(scheduler, NUM_STEPS, minimax_lora.PDD)
+    stub = types.SimpleNamespace(step_index=None)
+    head = network_pdd.ParallelHead(base, weight, bias, 1.0, lambda: stub, BLOCK, intervals)
+    reference = ReferenceHead(weight, bias)
+    step_sizes = reference_time_grid(12.0, NUM_STEPS).diff()
+    x = torch.randn(2, 5, HIDDEN)
+    for index in range(NUM_STEPS // BLOCK):
+        stub.step_index = None if index == 0 else index
+        reference.plan = reference_plan(step_sizes, index * BLOCK, BLOCK).float()
+        out = head(x)
+        ref = reference(x)
+        assert torch.allclose(out, ref, rtol=1e-4, atol=1e-4), f'block={index} maxdiff={(out - ref).abs().max().item()}' # float32 reduction order differs between tensordot and einsum
+        assert head.fused_index == index
+
+
+def test_parallel_head_strength_blend():
+    heads = make_heads()
+    weight, bias = heads.heads['proj_out']
+    base = torch.nn.Linear(HIDDEN, VIDEO_OUT)
+    intervals = network_pdd.grid_intervals(diffusers.MiniMaxH3Scheduler(shift=12.0), NUM_STEPS, minimax_lora.PDD)
+    stub = types.SimpleNamespace(step_index=3)
+    x = torch.randn(3, HIDDEN)
+    full = network_pdd.ParallelHead(base, weight, bias, 1.0, lambda: stub, BLOCK, intervals)(x)
+    off = network_pdd.ParallelHead(base, weight, bias, 0.0, lambda: stub, BLOCK, intervals)(x)
+    half = network_pdd.ParallelHead(base, weight, bias, 0.5, lambda: stub, BLOCK, intervals)(x)
+    assert torch.allclose(off, base(x), atol=1e-6), 'strength 0 is not the base projection'
+    assert torch.allclose(half, 0.5 * (full + base(x)), atol=1e-5), 'strength 0.5 is not the midpoint'
+
+
+def test_parallel_head_clamps_overflow():
+    heads = make_heads()
+    weight, bias = heads.heads['proj_out']
+    intervals = network_pdd.grid_intervals(diffusers.MiniMaxH3Scheduler(shift=12.0), NUM_STEPS, minimax_lora.PDD)
+    stub = types.SimpleNamespace(step_index=11)
+    head = network_pdd.ParallelHead(torch.nn.Linear(HIDDEN, VIDEO_OUT), weight, bias, 1.0, lambda: stub, BLOCK, intervals)
+    head(torch.randn(1, HIDDEN))
+    assert head.fused_index == 7 and head.overflow_warned
+
+
+def test_parallel_head_keeps_base_out_of_tree():
+    heads = make_heads()
+    weight, bias = heads.heads['proj_out']
+    intervals = network_pdd.grid_intervals(diffusers.MiniMaxH3Scheduler(shift=12.0), NUM_STEPS, minimax_lora.PDD)
+    head = network_pdd.ParallelHead(torch.nn.Linear(HIDDEN, VIDEO_OUT), weight, bias, 1.0, lambda: None, BLOCK, intervals)
+    assert set(dict(head.named_parameters())) == {'weight', 'bias'}, list(dict(head.named_parameters()))
+    assert len(list(head.children())) == 0
+    assert next(head.parameters()).dtype == torch.float32
+
+
+# ============================================================
+# Tests: install, restore, reconcile, pin
+# ============================================================
+
+def test_install_and_restore_round_trip():
+    pipe = StubPipe()
+    original_video, original_audio = pipe.transformer.proj_out, pipe.transformer.audio_proj_out
+    heads = make_heads()
+    assert network_pdd.install(pipe, make_net(heads), heads, minimax_lora.PDD, ['unet', 'transformer']) is True
+    assert isinstance(pipe.transformer.proj_out, network_pdd.ParallelHead)
+    assert isinstance(pipe.transformer.audio_proj_out, network_pdd.ParallelHead)
+    assert pipe.sdnext_pdd.steps == 9 and pipe.sdnext_pdd.component == 'transformer'
+    pipe.audio_scheduler.set_shift(3.0)
+    pipe.scheduler.set_timesteps(9)
+    pipe.audio_scheduler.set_timesteps(9)
+    out = pipe.transformer.audio_proj_out(torch.randn(2, HIDDEN))
+    assert out.shape == (2, AUDIO_OUT)
+    assert network_pdd.restore(pipe) is True
+    assert pipe.transformer.proj_out is original_video and pipe.transformer.audio_proj_out is original_audio
+    assert not hasattr(pipe, 'sdnext_pdd')
+    assert network_pdd.restore(pipe) is False
+
+
+def test_install_refuses_shape_mismatch():
+    pipe = StubPipe()
+    original = pipe.transformer.proj_out
+    heads = make_heads()
+    heads.heads['audio_proj_out'] = (torch.randn(NUM_STEPS, AUDIO_OUT + 1, HIDDEN), None)
+    assert network_pdd.install(pipe, make_net(heads), heads, minimax_lora.PDD, ['transformer']) is False
+    assert pipe.transformer.proj_out is original, 'partial install left a head behind'
+    assert not hasattr(pipe, 'sdnext_pdd')
+
+
+def test_install_needs_an_owner():
+    pipe = StubPipe()
+    heads = make_heads()
+    heads.heads['norm_out.linear'] = (torch.randn(NUM_STEPS, VIDEO_OUT, HIDDEN), None)
+    assert network_pdd.install(pipe, make_net(heads), heads, minimax_lora.PDD, ['transformer']) is False
+
+
+def test_reconcile_follows_loaded_set():
+    pipe = StubPipe()
+    heads = make_heads()
+    net = make_net(heads)
+    saved = network_pdd.arch_spec
+    network_pdd.arch_spec = lambda: minimax_lora.PDD
+    try:
+        assert network_pdd.reconcile(pipe, [net], ['transformer']) is True
+        assert network_pdd.reconcile(pipe, [net], ['transformer']) is False, 'unchanged set reinstalled'
+        stronger = make_net(heads, strength=0.5)
+        assert network_pdd.reconcile(pipe, [stronger], ['transformer']) is True, 'strength change not applied'
+        assert pipe.sdnext_pdd.strength == 0.5
+        assert network_pdd.reconcile(pipe, [types.SimpleNamespace(name='plain', te_multiplier=1.0, extras={})], ['transformer']) is True
+        assert not hasattr(pipe, 'sdnext_pdd')
+        assert network_pdd.reconcile(pipe, [], ['transformer']) is False
+    finally:
+        network_pdd.arch_spec = saved
+
+
+def test_pin_overrides_steps_and_shift():
+    pipe = StubPipe()
+    heads = make_heads()
+    assert network_pdd.install(pipe, make_net(heads), heads, minimax_lora.PDD, ['transformer']) is True
+    pipe.scheduler.set_shift(4.0)
+    pipe.audio_scheduler.set_shift(2.0)
+    p = types.SimpleNamespace(steps=30, task_args={'num_inference_steps': 30})
+    assert network_pdd.pin(p, pipe) == 9
+    assert p.steps == 9 and p.task_args['num_inference_steps'] == 9
+    assert pipe.num_timesteps == 8
+    assert pipe.scheduler.shift == 12.0 and pipe.audio_scheduler.shift == 3.0
+    network_pdd.restore(pipe)
+    assert network_pdd.pin(p, pipe) is None
+
+
+# ============================================================
+# Main
+# ============================================================
+
+def main():
+    cat = category('detect')
+    for fn in (test_detect_grid, test_detect_rejects_bad_block, test_load_collects_heads, test_load_without_metadata_is_none, test_load_without_heads_is_none):
+        run_test(cat, fn)
+    cat = category('math')
+    for fn in (test_grid_intervals_match_reference, test_steps_for_counts_terminal_sigma, test_parallel_head_matches_reference_per_block, test_parallel_head_strength_blend, test_parallel_head_clamps_overflow, test_parallel_head_keeps_base_out_of_tree):
+        run_test(cat, fn)
+    cat = category('lifecycle')
+    for fn in (test_install_and_restore_round_trip, test_install_refuses_shape_mismatch, test_install_needs_an_owner, test_reconcile_follows_loaded_set, test_pin_overrides_steps_and_shift):
+        run_test(cat, fn)
+    failed = sum(r['failed'] for r in results.values())
+    passed = sum(r['passed'] for r in results.values())
+    log.info(f'PDD tests: passed={passed} failed={failed}')
+    return 1 if failed else 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())