diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index fd0ff8730..01c130c95 100644 --- a/modules/lora/native_adapter.py +++ b/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') diff --git a/modules/lora/network.py b/modules/lora/network.py index 899e6940c..5f23407ca 100644 --- a/modules/lora/network.py +++ b/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 diff --git a/modules/lora/network_pdd.py b/modules/lora/network_pdd.py new file mode 100644 index 000000000..82cf7f35a --- /dev/null +++ b/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 diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 87dde344f..cab987fdf 100644 --- a/modules/lora/networks.py +++ b/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") diff --git a/modules/processing_args.py b/modules/processing_args.py index 4ac2781ff..2833f4f82 100644 --- a/modules/processing_args.py +++ b/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) diff --git a/pipelines/minimax/minimax_lora.py b/pipelines/minimax/minimax_lora.py index 26961b684..171f47d42 100644 --- a/pipelines/minimax/minimax_lora.py +++ b/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, ), ) diff --git a/test/test-pdd.py b/test/test-pdd.py new file mode 100644 index 000000000..7eab00d7b --- /dev/null +++ b/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())