mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-10-11 22:49:59 +00:00
feat(lora): pdd exports with a trained sigma grid and replaced weights
Inference-only PDD exports such as alibaba-pai's Qwen-Image-2.1-Fun-Acc-4Step carry no pdd_num_steps: the grid is pdd_sigmas in the metadata or in the pdd_config.json the release keeps beside the weights, and the file replaces norm weights outright. Those parameters are set with the heads, blended toward the loaded value by strength and restored on removal. The pin passes the grid's block boundaries as sigmas and turns off dynamic shifting, static shift and terminal stretch on that generation's scheduler, which would otherwise reshape them, and leaves num_timesteps alone where the pipeline exposes it as a read-only property. Qwen-Image 2.1 gets an ArchSpec and the PDD loader in its chain; a file with heads but no grid warns and loads its backbone only.
This commit is contained in:
CalamitousFelicitousness
@@ -56,6 +56,7 @@ Plus quite a few other improvements and inevitable bug-fixes, see full changelog
|
||||
- edit-models: simplify usage of multi-image reference editing
|
||||
automatically handle: flux.2, glm-image, nano-banana, joy-image-edit, qwen-image-edit
|
||||
- lora: native loading for qwen-image 2.1, including the fused mlp layout ai-toolkit and musubi-tuner train against
|
||||
- lora: alibaba-pai qwen-image 2.1 fun-acc 4-step runs on the sigma grid it was distilled for
|
||||
- processors: `DWPose`, `RTMW`, `RTMO`: image processor updates, thanks @kirtasshh
|
||||
refactored to use `ONNX` instead of obsolete `mmpose` lib and with additional configurable settings
|
||||
- **UI**
|
||||
|
||||
+170
-41
@@ -4,10 +4,15 @@ A PDD file pairs a backbone LoRA with the output projections repeated once per i
|
||||
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.
|
||||
|
||||
Inference-only exports can also carry the grid itself as ``pdd_sigmas`` (or in the ``pdd_config.json`` the
|
||||
release keeps beside its weights) and parameters that replace the loaded values outright, such as fine-tuned
|
||||
norm weights; the pin then hands the scheduler that exact grid.
|
||||
"""
|
||||
|
||||
import os
|
||||
import copy
|
||||
import json
|
||||
import weakref
|
||||
import torch
|
||||
from modules.logger import log
|
||||
@@ -18,6 +23,9 @@ debug_log = log.trace if os.environ.get('SD_LORA_DEBUG', None) is not None else
|
||||
|
||||
METADATA_STEPS = 'pdd_num_steps'
|
||||
METADATA_BLOCK = 'pdd_block_size'
|
||||
METADATA_SIGMAS = 'pdd_sigmas'
|
||||
GRID_KEYS = (METADATA_STEPS, METADATA_BLOCK, METADATA_SIGMAS)
|
||||
SIDECAR = 'pdd_config.json'
|
||||
EXTRAS_KEY = 'pdd'
|
||||
|
||||
|
||||
@@ -35,46 +43,73 @@ class ArchSpec:
|
||||
|
||||
|
||||
class ParallelHeads:
|
||||
"""The head tensors of one file and the grid they were trained on."""
|
||||
"""The head tensors of one file, the grid they were trained on and the parameters the file replaces outright."""
|
||||
|
||||
def __init__(self, num_steps, block_size, heads):
|
||||
def __init__(self, num_steps, block_size, heads, sigmas=None, replaced=None):
|
||||
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.sigmas = sigmas # the training grid descending from 1 to 0, or None to take it from the scheduler
|
||||
self.replaced = replaced or {} # parameter path -> the value the file sets
|
||||
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):
|
||||
def __init__(self, name, strength, component, modules, heads, spec, replaced=None):
|
||||
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.replaced = replaced or {} # parameter path -> (parameter, original data)
|
||||
self.steps = spec.steps_for(heads.nfe)
|
||||
|
||||
|
||||
def grid_sigmas(metadata):
|
||||
"""The explicit training grid under ``pdd_sigmas`` as float64, None when absent, ValueError unless it descends from 1 to 0."""
|
||||
raw = (metadata or {}).get(METADATA_SIGMAS)
|
||||
if raw is None:
|
||||
return None
|
||||
values = json.loads(raw) if isinstance(raw, str) else list(raw)
|
||||
sigmas = torch.tensor([float(value) for value in values], dtype=torch.float64)
|
||||
if sigmas.numel() < 2 or sigmas[0] != 1 or sigmas[-1] != 0 or not bool((sigmas.diff() < 0).all()):
|
||||
raise ValueError(f'sigmas={values}')
|
||||
return sigmas
|
||||
|
||||
|
||||
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:
|
||||
sigmas = grid_sigmas(metadata)
|
||||
if METADATA_STEPS not in metadata and sigmas is None:
|
||||
return None
|
||||
num_steps = int(metadata[METADATA_STEPS])
|
||||
num_steps = int(metadata[METADATA_STEPS]) if METADATA_STEPS in metadata else sigmas.numel() - 1
|
||||
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}')
|
||||
if sigmas is not None and sigmas.numel() != num_steps + 1:
|
||||
raise ValueError(f'grid={num_steps} sigmas={sigmas.numel()}')
|
||||
return num_steps, block_size
|
||||
|
||||
|
||||
def adapter_key(key):
|
||||
"""True for keys of an adapter family (LoRA factors, alphas, LyCORIS tensors) rather than plain parameters."""
|
||||
from modules.lora import native_adapter
|
||||
markers = (native_adapter.LORA_MARKERS + native_adapter.LOKR_MARKERS + native_adapter.LOHA_MARKERS + native_adapter.OFT_MARKERS
|
||||
+ native_adapter.IA3_MARKERS + native_adapter.GLORA_MARKERS + native_adapter.NORM_MARKERS + native_adapter.FULL_MARKERS + ('.alpha', '.dora_scale'))
|
||||
return any(marker in key for marker in markers)
|
||||
|
||||
|
||||
def load(name, metadata, state_dict):
|
||||
"""Collect the per-interval head tensors of a file, or None when the file carries no PDD grid."""
|
||||
"""Collect the per-interval heads and the replaced parameters 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')
|
||||
sigmas = grid_sigmas(metadata)
|
||||
except (TypeError, ValueError) as e:
|
||||
log.error(f'Network load: type=PDD name="{name}" {e} unusable grid')
|
||||
return None
|
||||
if grid is None:
|
||||
return None
|
||||
@@ -87,34 +122,68 @@ def load(name, metadata, state_dict):
|
||||
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)
|
||||
head_keys = {f'{path}.{leaf}' for path in heads for leaf in ('weight', 'bias')}
|
||||
replaced = {key: tensor for key, tensor in state_dict.items() if key not in head_keys and key.rsplit('.', 1)[-1] in ('weight', 'bias') and not adapter_key(key)}
|
||||
log.debug(f'Network load: type=PDD name="{name}" grid={num_steps} block={block_size} nfe={num_steps // block_size} heads={list(heads)} sigmas={None if sigmas is None else sigmas.tolist()} replaced={len(replaced)}')
|
||||
return ParallelHeads(num_steps, block_size, heads, sigmas=sigmas, replaced=replaced)
|
||||
|
||||
|
||||
def has_grid(metadata):
|
||||
return METADATA_STEPS in metadata or METADATA_SIGMAS in metadata
|
||||
|
||||
|
||||
def header_metadata(filename, name):
|
||||
"""The metadata read from the file itself, for when the cached metadata lacks the grid; warns when head-shaped tensors have no grid."""
|
||||
"""The metadata read from the file itself, for when the cached metadata lacks the grid."""
|
||||
from safetensors import safe_open
|
||||
try:
|
||||
with safe_open(filename, framework='pt', device='cpu') as f:
|
||||
metadata = f.metadata() or {}
|
||||
if METADATA_STEPS in metadata:
|
||||
if has_grid(metadata):
|
||||
log.debug(f'Network load: type=PDD name="{name}" grid read from the file header')
|
||||
else:
|
||||
heads = sum(1 for key in f.keys() if key.endswith('.weight') and len(f.get_slice(key).get_shape()) == 3)
|
||||
if heads > 0:
|
||||
log.warning(f'Network load: type=PDD name="{name}" heads={heads} no {METADATA_STEPS} metadata: heads ignored')
|
||||
return metadata
|
||||
except Exception as e:
|
||||
log.warning(f'Network load: type=PDD name="{name}" header {e}')
|
||||
return {}
|
||||
|
||||
|
||||
def sidecar_metadata(filename, name):
|
||||
"""The grid keys of the ``pdd_config.json`` a release keeps beside its weights, as metadata strings; {} without one."""
|
||||
path = os.path.join(os.path.dirname(filename), SIDECAR)
|
||||
if not os.path.isfile(path):
|
||||
return {}
|
||||
try:
|
||||
with open(path, encoding='utf8') as f:
|
||||
config = json.load(f)
|
||||
except (OSError, ValueError) as e:
|
||||
log.warning(f'Network load: type=PDD name="{name}" config="{path}" {e}')
|
||||
return {}
|
||||
metadata = {key: json.dumps(value) if isinstance(value, list) else str(value) for key, value in config.items() if key in GRID_KEYS}
|
||||
if has_grid(metadata):
|
||||
log.debug(f'Network load: type=PDD name="{name}" grid read from "{path}"')
|
||||
return metadata
|
||||
|
||||
|
||||
def headless_warning(filename, name):
|
||||
"""Warn when a file holds head-shaped tensors but no grid was found anywhere."""
|
||||
from safetensors import safe_open
|
||||
try:
|
||||
with safe_open(filename, framework='pt', device='cpu') as f:
|
||||
heads = sum(1 for key in f.keys() if key.endswith('.weight') and len(f.get_slice(key).get_shape()) == 3)
|
||||
except Exception:
|
||||
return
|
||||
if heads > 0:
|
||||
log.warning(f'Network load: type=PDD name="{name}" heads={heads} no {METADATA_STEPS} or {METADATA_SIGMAS} metadata and no {SIDECAR} beside the file: heads ignored')
|
||||
|
||||
|
||||
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:
|
||||
if not has_grid(metadata):
|
||||
metadata = header_metadata(network_on_disk.filename, name) # the metadata cache keeps a failed read forever and --no-metadata returns nothing
|
||||
if METADATA_STEPS not in metadata:
|
||||
if not has_grid(metadata):
|
||||
metadata = {**metadata, **sidecar_metadata(network_on_disk.filename, name)}
|
||||
if not has_grid(metadata):
|
||||
headless_warning(network_on_disk.filename, name)
|
||||
return None
|
||||
from modules.lora import native_adapter
|
||||
state_dict = native_adapter.read_state_dict(network_on_disk.filename, what='network')
|
||||
@@ -194,8 +263,10 @@ class ParallelHead(torch.nn.Module):
|
||||
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."""
|
||||
def grid_intervals(scheduler, num_steps, spec, sigmas=None):
|
||||
"""Interval lengths of the training grid in ascending time: from the file's own grid, or a pristine scheduler copy immune to a live shift override."""
|
||||
if sigmas is not None:
|
||||
return (1.0 - sigmas).diff()
|
||||
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)
|
||||
@@ -228,8 +299,37 @@ def target_shape(module):
|
||||
return tuple(weight.shape) if weight is not None else None
|
||||
|
||||
|
||||
def undo(modules, replaced):
|
||||
"""Put back swapped projections and replaced parameter values."""
|
||||
for parent, attr, original in modules.values():
|
||||
setattr(parent, attr, original)
|
||||
for parameter, original in replaced.values():
|
||||
parameter.data = original
|
||||
|
||||
|
||||
def replace_parameters(net, heads, component, strength):
|
||||
"""Set every parameter the file replaces, blended toward the loaded value by strength; None on a missing or misshapen target."""
|
||||
replaced = {}
|
||||
for path, value in heads.replaced.items():
|
||||
try:
|
||||
parameter = component.get_parameter(path)
|
||||
except AttributeError:
|
||||
parameter = None
|
||||
if parameter is None or tuple(parameter.shape) != tuple(value.shape):
|
||||
log.error(f'Network load: type=PDD name="{net.name}" parameter={path} shape={list(value.shape)} module={None if parameter is None else list(parameter.shape)} no matching parameter')
|
||||
undo({}, replaced)
|
||||
return None
|
||||
original = parameter.data
|
||||
target = value.to(device=original.device, dtype=torch.float32)
|
||||
if strength != 1.0:
|
||||
target = original.to(dtype=torch.float32) + strength * (target - original.to(dtype=torch.float32))
|
||||
parameter.data = target.to(dtype=original.dtype)
|
||||
replaced[path] = (parameter, original)
|
||||
return replaced
|
||||
|
||||
|
||||
def install(pipe, net, heads, spec, components):
|
||||
"""Swap the heads into the component that owns their projections; True when the module tree changed."""
|
||||
"""Swap the heads into the component that owns their projections and set the parameters the file replaces; 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')
|
||||
@@ -242,16 +342,14 @@ def install(pipe, net, heads, spec, components):
|
||||
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)
|
||||
undo(modules, {})
|
||||
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
|
||||
intervals = grid_intervals(scheduler, heads.num_steps, spec, heads.sigmas) if scheduler is not None or heads.sigmas 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)
|
||||
undo(modules, {})
|
||||
return False
|
||||
def get_scheduler(name=scheduler_name):
|
||||
owner_pipe = pipe_ref()
|
||||
@@ -263,20 +361,23 @@ def install(pipe, net, heads, spec, components):
|
||||
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}')
|
||||
replaced = replace_parameters(net, heads, component, strength)
|
||||
if replaced is None:
|
||||
undo(modules, {})
|
||||
return False
|
||||
pipe.sdnext_pdd = Installed(net.name, strength, component_name, modules, heads, spec, replaced)
|
||||
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} replaced={len(replaced)} strength={strength}')
|
||||
return True
|
||||
|
||||
|
||||
def restore(pipe):
|
||||
"""Put the original projections back; True when heads were installed."""
|
||||
"""Put the original projections and parameter values 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)
|
||||
undo(state.modules, state.replaced)
|
||||
del pipe.sdnext_pdd
|
||||
log.info(f'Network unload: type=PDD name="{state.name}" component={state.component} heads={list(state.modules)}')
|
||||
log.info(f'Network unload: type=PDD name="{state.name}" component={state.component} heads={list(state.modules)} replaced={len(state.replaced)}')
|
||||
return True
|
||||
|
||||
|
||||
@@ -311,26 +412,54 @@ def reconcile(pipe, loaded, components):
|
||||
return install(pipe, net, net.extras[EXTRAS_KEY], spec, components) or changed
|
||||
|
||||
|
||||
def settable(obj, name):
|
||||
"""False for a property without a setter."""
|
||||
attribute = getattr(type(obj), name, None)
|
||||
return not isinstance(attribute, property) or attribute.fset is not None
|
||||
|
||||
|
||||
SIGMA_TRANSFORMS = {'use_dynamic_shifting': False, 'shift_terminal': None, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'invert_sigmas': False}
|
||||
|
||||
|
||||
def take_sigmas_as_given(scheduler):
|
||||
"""Turn off every transform a flow scheduler applies to explicit sigmas; the sampler builds a fresh scheduler for each generation."""
|
||||
config = getattr(scheduler, 'config', None) or {}
|
||||
flags = {key: value for key, value in SIGMA_TRANSFORMS.items() if key in config}
|
||||
if flags and hasattr(scheduler, 'register_to_config'):
|
||||
scheduler.register_to_config(**flags)
|
||||
if hasattr(scheduler, 'set_shift'):
|
||||
scheduler.set_shift(1.0)
|
||||
|
||||
|
||||
def pin(p, model):
|
||||
"""Hold a generation on the distilled evaluation count and shipped schedule while heads are installed; returns the scheduler step argument 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']
|
||||
sigmas = None
|
||||
if state.heads.sigmas is not None:
|
||||
scheduler = getattr(model, state.spec.default_scheduler, None)
|
||||
if scheduler is not None:
|
||||
take_sigmas_as_given(scheduler)
|
||||
sigmas = state.heads.sigmas[::state.heads.block_size][:-1].tolist() # block boundaries; the scheduler appends the terminal zero
|
||||
if getattr(p, 'task_args', None) is not None:
|
||||
p.task_args['sigmas'] = sigmas
|
||||
else:
|
||||
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.heads.nfe # what the user-facing step count means: transformer evaluations
|
||||
if getattr(p, 'task_args', None) is not None:
|
||||
p.task_args['num_inference_steps'] = state.steps # the scheduler argument that yields that many grid intervals
|
||||
if getattr(model, 'num_timesteps', None) is not None:
|
||||
model.num_timesteps = state.heads.nfe # the progress total counts transformer evaluations
|
||||
if getattr(model, 'num_timesteps', None) is not None and settable(model, 'num_timesteps'):
|
||||
model.num_timesteps = state.heads.nfe # the progress total counts transformer evaluations; a read-only property is the pipeline counting its own timesteps
|
||||
extra = getattr(p, 'extra_generation_params', None)
|
||||
if extra is not None:
|
||||
extra.update({state.spec.shift_keys[name]: shift for name, shift in shifts.items() if name in state.spec.shift_keys})
|
||||
log.info(f'Network: type=PDD name="{state.name}" steps={state.heads.nfe} requested={requested} grid_steps={state.steps} shift={shifts}')
|
||||
log.info(f'Network: type=PDD name="{state.name}" steps={state.heads.nfe} requested={requested} grid_steps={state.steps} shift={shifts} sigmas={sigmas}')
|
||||
return state.steps
|
||||
|
||||
@@ -15,12 +15,19 @@ under every prefix. ComfyUI checkpoints fuse the SwiGLU input as
|
||||
``img_mlp.gate_up`` with rows ``[gate_layer; proj]``, and so do the LoRAs
|
||||
ai-toolkit and musubi-tuner train against them; :func:`resolve_targets` splits
|
||||
that projection into the two diffusers modules.
|
||||
|
||||
Parallel decoding exports (alibaba-pai Fun-Acc) carry their own sigma grid, so
|
||||
:data:`PDD` needs no scheduler mapping: the pin hands the pipeline that grid.
|
||||
"""
|
||||
|
||||
from modules.lora import native_adapter
|
||||
from modules.lora import native_adapter, network_pdd
|
||||
from modules.lora.native_adapter import ChunkSpec
|
||||
|
||||
|
||||
# One output projection on the pipeline scheduler; with explicit sigmas the step count is the interval count.
|
||||
PDD = network_pdd.ArchSpec()
|
||||
|
||||
|
||||
# === Arch-specific prefix configuration ===
|
||||
|
||||
KNOWN_PREFIXES = native_adapter.KNOWN_PREFIXES_DEFAULT
|
||||
@@ -142,5 +149,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,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,6 +19,7 @@ Usage:
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import types
|
||||
|
||||
import torch
|
||||
@@ -46,6 +47,7 @@ from modules import shared # pylint: disable=wrong-import-position,unused-impor
|
||||
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
|
||||
from pipelines.qwen import qwen21_lora # pylint: disable=wrong-import-position
|
||||
|
||||
|
||||
NUM_STEPS = 32
|
||||
@@ -152,6 +154,45 @@ def make_net(heads, strength=1.0):
|
||||
return types.SimpleNamespace(name='pdd-test', te_multiplier=strength, extras={network_pdd.EXTRAS_KEY: heads})
|
||||
|
||||
|
||||
# Qwen-Image-2.1-Fun-Acc-4Step: an inference-only export with one head per step, its own sigma grid and replaced norm weights.
|
||||
QWEN_SIGMAS = [1.0, 0.9169867038726807, 0.7861579060554504, 0.5494909882545471, 0.0]
|
||||
QWEN_SCHEDULER = dict(use_dynamic_shifting=True, base_shift=0.5, max_shift=0.9, base_image_seq_len=256, max_image_seq_len=8192, shift_terminal=0.02, time_shift_type='exponential')
|
||||
|
||||
|
||||
class StubQwenPipe:
|
||||
"""A single-stream model: one bias-free output projection, a norm the export replaces, a dynamic-shift flow scheduler,
|
||||
and the read-only num_timesteps property diffusers pipelines expose."""
|
||||
|
||||
def __init__(self):
|
||||
self.transformer = torch.nn.Module()
|
||||
self.transformer.proj_out = torch.nn.Linear(HIDDEN, VIDEO_OUT, bias=False)
|
||||
self.transformer.block = torch.nn.Module()
|
||||
self.transformer.block.norm_q = torch.nn.RMSNorm(8)
|
||||
self.scheduler = diffusers.FlowMatchEulerDiscreteScheduler(**QWEN_SCHEDULER)
|
||||
self._num_timesteps = 40
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
|
||||
def qwen_state_dict():
|
||||
torch.manual_seed(1)
|
||||
return {
|
||||
'proj_out.weight': torch.randn(4, VIDEO_OUT, HIDDEN),
|
||||
'block.norm_q.weight': torch.full((8,), 2.0),
|
||||
'block.attn.to_q.lora_down': torch.zeros(4, HIDDEN),
|
||||
'block.attn.to_q.lora_up': torch.zeros(HIDDEN, 4),
|
||||
'block.attn.to_k.lora_A.weight': torch.zeros(4, HIDDEN),
|
||||
'block.attn.to_k.lora_B.weight': torch.zeros(HIDDEN, 4),
|
||||
'block.attn.to_k.alpha': torch.tensor(4.0),
|
||||
}
|
||||
|
||||
|
||||
def qwen_heads():
|
||||
return network_pdd.load('acc', {'pdd_sigmas': json.dumps(QWEN_SIGMAS)}, qwen_state_dict())
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Tests: detection and loading
|
||||
# ============================================================
|
||||
@@ -369,6 +410,113 @@ def test_pin_overrides_steps_and_shift():
|
||||
assert network_pdd.pin(p, pipe) is None
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Tests: explicit grids and replaced parameters
|
||||
# ============================================================
|
||||
|
||||
def test_detect_grid_from_sigmas():
|
||||
metadata = {'pdd_sigmas': json.dumps(QWEN_SIGMAS)}
|
||||
assert network_pdd.detect(metadata) == (4, 1)
|
||||
sigmas = network_pdd.grid_sigmas(metadata)
|
||||
assert sigmas.dtype == torch.float64 and sigmas.tolist() == QWEN_SIGMAS
|
||||
assert network_pdd.detect({'pdd_num_steps': '4', 'pdd_sigmas': json.dumps(QWEN_SIGMAS)}) == (4, 1)
|
||||
|
||||
|
||||
def test_detect_rejects_bad_sigmas():
|
||||
for sigmas in ([1.0, 0.5], [0.9, 0.5, 0.0], [1.0, 0.5, 0.6, 0.0]):
|
||||
try:
|
||||
network_pdd.detect({'pdd_sigmas': json.dumps(sigmas)})
|
||||
except ValueError:
|
||||
continue
|
||||
raise AssertionError(f'sigmas {sigmas} accepted')
|
||||
try:
|
||||
network_pdd.detect({'pdd_num_steps': '8', 'pdd_sigmas': json.dumps(QWEN_SIGMAS)})
|
||||
except ValueError:
|
||||
return True
|
||||
raise AssertionError('an 8-step grid accepted 5 sigmas')
|
||||
|
||||
|
||||
def test_load_collects_replaced_parameters():
|
||||
heads = qwen_heads()
|
||||
assert heads is not None and set(heads.heads) == {'proj_out'} and heads.heads['proj_out'][1] is None
|
||||
assert set(heads.replaced) == {'block.norm_q.weight'}, f'replaced={list(heads.replaced)}'
|
||||
assert heads.nfe == 4 and heads.sigmas.tolist() == QWEN_SIGMAS
|
||||
|
||||
|
||||
def test_try_load_reads_the_sidecar_config():
|
||||
import tempfile
|
||||
from safetensors.torch import save_file
|
||||
with tempfile.TemporaryDirectory() as folder:
|
||||
path = os.path.join(folder, 'acc.safetensors')
|
||||
save_file(qwen_state_dict(), path, metadata={'format': 'qwenimage21_extracted_prefused_v1'})
|
||||
disk = types.SimpleNamespace(name='acc', filename=path, metadata={})
|
||||
assert network_pdd.try_load('acc', disk, 1.0) is None, 'no grid anywhere must leave the heads out'
|
||||
with open(os.path.join(folder, network_pdd.SIDECAR), 'w', encoding='utf8') as f:
|
||||
json.dump({'pdd_num_steps': 4, 'pdd_block_size': 1, 'pdd_sigmas': QWEN_SIGMAS, 'lora_rank': 64}, f)
|
||||
net = network_pdd.try_load('acc', disk, 1.0)
|
||||
assert net is not None, 'the sidecar grid must be enough'
|
||||
heads = net.extras[network_pdd.EXTRAS_KEY]
|
||||
assert heads.nfe == 4 and heads.sigmas.tolist() == QWEN_SIGMAS
|
||||
|
||||
|
||||
def test_install_replaces_and_restores_parameters():
|
||||
pipe = StubQwenPipe()
|
||||
original = pipe.transformer.block.norm_q.weight.detach().clone()
|
||||
heads = qwen_heads()
|
||||
assert network_pdd.install(pipe, make_net(heads), heads, qwen21_lora.PDD, ['transformer']) is True
|
||||
assert torch.equal(pipe.transformer.block.norm_q.weight, torch.full((8,), 2.0))
|
||||
network_pdd.restore(pipe)
|
||||
assert torch.equal(pipe.transformer.block.norm_q.weight, original)
|
||||
assert network_pdd.install(pipe, make_net(heads, strength=0.5), heads, qwen21_lora.PDD, ['transformer']) is True
|
||||
assert torch.allclose(pipe.transformer.block.norm_q.weight, original + 0.5 * (2.0 - original))
|
||||
network_pdd.restore(pipe)
|
||||
assert torch.equal(pipe.transformer.block.norm_q.weight, original)
|
||||
|
||||
|
||||
def test_install_refuses_misshapen_replacement():
|
||||
pipe = StubQwenPipe()
|
||||
original_head = pipe.transformer.proj_out
|
||||
heads = qwen_heads()
|
||||
heads.replaced['block.norm_q.weight'] = torch.ones(16)
|
||||
assert network_pdd.install(pipe, make_net(heads), heads, qwen21_lora.PDD, ['transformer']) is False
|
||||
assert pipe.transformer.proj_out is original_head and not hasattr(pipe, 'sdnext_pdd')
|
||||
|
||||
|
||||
def test_one_head_per_step_on_a_block1_grid():
|
||||
pipe = StubQwenPipe()
|
||||
heads = qwen_heads()
|
||||
assert network_pdd.install(pipe, make_net(heads), heads, qwen21_lora.PDD, ['transformer']) is True
|
||||
x = torch.randn(3, HIDDEN)
|
||||
for index in range(4):
|
||||
pipe.scheduler._step_index = index # pylint: disable=protected-access
|
||||
expected = torch.nn.functional.linear(x, heads.heads['proj_out'][0][index])
|
||||
assert torch.allclose(pipe.transformer.proj_out(x), expected, atol=1e-6), f'step {index} did not select its own head'
|
||||
network_pdd.restore(pipe)
|
||||
|
||||
|
||||
def test_pin_hands_the_exact_grid():
|
||||
pipe = StubQwenPipe()
|
||||
try:
|
||||
pipe.num_timesteps = 4
|
||||
except AttributeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError('the stub accepts num_timesteps assignment; it does not model the diffusers property')
|
||||
heads = qwen_heads()
|
||||
assert network_pdd.install(pipe, make_net(heads), heads, qwen21_lora.PDD, ['transformer']) is True
|
||||
p = types.SimpleNamespace(steps=40, task_args={}, extra_generation_params={})
|
||||
assert network_pdd.pin(p, pipe) == 4 and p.steps == 4 and p.task_args['num_inference_steps'] == 4
|
||||
assert pipe.num_timesteps == 40, 'the pipeline counts its own timesteps'
|
||||
assert p.task_args['sigmas'] == QWEN_SIGMAS[:-1]
|
||||
pipe.scheduler.set_timesteps(sigmas=p.task_args['sigmas'], mu=1.31) # the pipeline passes a resolution-dependent mu
|
||||
grid = torch.tensor(QWEN_SIGMAS, dtype=torch.float64)
|
||||
assert torch.allclose(pipe.scheduler.sigmas.double(), grid, atol=1e-6), f'sigmas={pipe.scheduler.sigmas.tolist()}'
|
||||
unpinned = diffusers.FlowMatchEulerDiscreteScheduler(**QWEN_SCHEDULER)
|
||||
unpinned.set_timesteps(sigmas=p.task_args['sigmas'], mu=1.31)
|
||||
assert not torch.allclose(unpinned.sigmas.double(), grid, atol=1e-3), 'the unpinned scheduler also lands on the grid; the check does not discriminate'
|
||||
network_pdd.restore(pipe)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
@@ -383,6 +531,10 @@ def main():
|
||||
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)
|
||||
cat = category('grid')
|
||||
for fn in (test_detect_grid_from_sigmas, test_detect_rejects_bad_sigmas, test_load_collects_replaced_parameters, test_try_load_reads_the_sidecar_config,
|
||||
test_install_replaces_and_restores_parameters, test_install_refuses_misshapen_replacement, test_one_head_per_step_on_a_block1_grid, test_pin_hands_the_exact_grid):
|
||||
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}')
|
||||
|
||||
Reference in New Issue
Block a user