diffusers-workflow 0.4.0__py3-none-any.whl

dw/assessment_rules.py

@@ -0,0 +1,159 @@
+"""The rules the assessment probes read their measurements against (#387).
+
+A probe (`dw/tasks/assess.py`) measures a finished file and reports every
+number it took; a rule here names one of those numbers and the threshold
+past which it is worth a look, and a crossing becomes a `finding` in the
+probe's answer. Findings are places to look, not verdicts: nothing in the
+engine acts on one, no run fails for one, and a finding a person has looked
+at and accepted is simply left alone. That is the authority rule, and it is
+why the thresholds sit in one table - a number an agent is told to trust
+has to be one somebody can find, argue with and change in one place.
+
+Each entry names the probe that reports the field, the field (a key of the
+probe's per-shot or per-seam record, or of its answer itself -
+`tests/test_assessment_rules.py` pins every one to a real probe field so a
+rename cannot leave a rule reading nothing), how the value is compared,
+the threshold, the severity and what a crossing says. `magnitude` compares
+the value's absolute size, for a signed measurement whose direction is not
+the problem. `shot_level_spread` shares its threshold with the run-time
+`LEVEL_SPREAD_WARN_DB`, so the join-time warning and the after-the-fact
+probe cannot disagree about the same cut.
+
+Two rules carry a guard the table alone cannot express, applied by the
+probe and named here in `unless` so it is written down beside the number:
+`seam_hole` holds only while both sides of the seam are voiced, and
+`seam_frame_jump` does not fire at a seam whose incoming shot is marked
+`hard_cut: true` - a cut meant as a cut.
+
+The thresholds were settled against real runs on the server in stage D
+(#386); `docs/WORKFLOW_GUIDE.md` and `docs/TASKS.md` quote every one, and
+`tests/test_assessment_rules.py` pins those quotes to this table.
+"""
+
+from .tasks.audio_utils import LEVEL_SPREAD_WARN_DB
+
+SEVERITIES = ("info", "warn")
+
+# How loud both sides of a seam have to be for a quiet join to be a hole
+# rather than a pause the shots themselves hold
+HOLE_VOICED_DBFS = -30.0
+
+COMPARATORS = {
+    ">": lambda value, threshold: value > threshold,
+    ">=": lambda value, threshold: value >= threshold,
+    "<": lambda value, threshold: value < threshold,
+}
+
+RULES = (
+    {
+        "name": "shot_level_spread",
+        "probe": "analyze_shots",
+        "field": "rms_range_db",
+        "comparator": ">=",
+        "threshold": LEVEL_SPREAD_WARN_DB,
+        "severity": "warn",
+        "says": "the shots sit this far apart in level - a jump a listener hears at the cut",
+    },
+    {
+        "name": "seam_level_step",
+        "probe": "analyze_seams",
+        "field": "level_step_db",
+        "comparator": ">",
+        "threshold": 3.0,
+        "severity": "warn",
+        "says": "the shots either side of the seam sit this far apart in level",
+    },
+    {
+        "name": "seam_click",
+        "probe": "analyze_seams",
+        "field": "click_db",
+        "comparator": ">",
+        "threshold": 12.0,
+        "severity": "warn",
+        "says": "the join peaks this far above the audio either side of it - an audible click",
+    },
+    {
+        "name": "seam_hole",
+        "probe": "analyze_seams",
+        "field": "floor_dbfs",
+        "comparator": "<",
+        "threshold": -50.0,
+        "severity": "warn",
+        "says": "the track drops out at the join while both sides are voiced",
+        "unless": f"either side's rms is at or below {HOLE_VOICED_DBFS} dBFS",
+    },
+    {
+        "name": "seam_frame_jump",
+        "probe": "analyze_seams",
+        "field": "jump_ratio",
+        "comparator": ">",
+        # Above every intended cut the stage D field runs measured (#386):
+        # seven ordinary cuts between generated shots read 9.9-20.3, dissolves
+        # about 1.8. At the original 8 the rule fired on every cut
+        "threshold": 25.0,
+        "severity": "info",
+        "says": "the picture changes this many times more across the seam than inside either shot",
+        "unless": "the incoming shot is marked hard_cut: true",
+    },
+    {
+        "name": "sync_drift",
+        "probe": "analyze_sync_drift",
+        "field": "end_offset_ms",
+        "comparator": ">",
+        "threshold": 40.0,
+        "severity": "warn",
+        "magnitude": True,
+        "says": "by this shot's end the audio sits this far off the picture",
+    },
+    {
+        "name": "sync_length",
+        "probe": "analyze_sync_drift",
+        "field": "length_delta_ms",
+        "comparator": ">",
+        "threshold": 40.0,
+        "severity": "warn",
+        "magnitude": True,
+        "says": "the soundtrack and the picture differ in length by this much",
+    },
+)
+
+RULES_BY_NAME = {rule["name"]: rule for rule in RULES}
+
+
+def rules_for(probe):
+    """The rules one probe's measurements are read against, in table order."""
+    return [rule for rule in RULES if rule["probe"] == probe]
+
+
+def crosses(rule, value):
+    """Whether a measured value crosses a rule's threshold. None never does -
+    a measurement that could not be taken (a silent window, no soundtrack)
+    is not a finding."""
+    if value is None:
+        return False
+    if rule.get("magnitude"):
+        value = abs(value)
+    return COMPARATORS[rule["comparator"]](value, rule["threshold"])
+
+
+def finding(rule, value, at):
+    """A crossing as the probe reports it."""
+    return {
+        "rule": rule["name"],
+        "severity": rule["severity"],
+        "at": at,
+        "value": value,
+        "threshold": rule["threshold"],
+        "says": rule["says"],
+    }
+
+
+__all__ = [
+    "HOLE_VOICED_DBFS",
+    "RULES",
+    "RULES_BY_NAME",
+    "SEVERITIES",
+    "crosses",
+    "finding",
+    "rules_for",
+]

dw/assets.py

@@ -0,0 +1,130 @@
+"""The asset library: input media a workflow references by name.
+
+A workflow's media paths resolve against the workflow file's own directory,
+which means a workflow that reads anything has to keep that thing beside it -
+the reason generated media ends up gitignored inside a source tree. An
+'asset:name' reference is rooted at the asset library instead, the way
+'prompt:name' is rooted at the prompt library, so the same reference means the
+same file from every workflow and neither has to live next to the other.
+
+A reference resolves to a path, not to a value: 'asset:frames/iris.jpg'
+becomes the absolute path of that file, and whatever would have loaded a path
+written there loads it unchanged.
+"""
+
+import contextvars
+import logging
+import os
+
+from .security import validate_asset_reference, validate_path
+from .workspace import ASSETS_SUBDIR, discover_library, library_fallbacks
+
+logger = logging.getLogger("dw")
+
+# The prefix marking a value as a reference to a stored asset
+ASSET_PREFIX = "asset:"
+
+# Set by an entry point from --asset-dir, and inherited by a spawned worker,
+# the way DW_PROMPT_DIR is
+ASSET_DIR_ENV_VAR = "DW_ASSET_DIR"
+
+
+# The asset library of the run in progress. A server holds several
+# workspaces and each has its own assets, so this cannot be a process-wide
+# environment variable there the way the prompt library can - there is one
+# prompt library, shared, but assets belong to a workspace. Set per job by
+# the worker; unset for the CLI and REPL, which have one workspace per
+# process and read the environment below
+_active_asset_dir = contextvars.ContextVar("dw_asset_dir", default=None)
+
+
+def activate_asset_dir(directory):
+    """Make an asset library the active one; returns a token for deactivate."""
+    return _active_asset_dir.set(directory)
+
+
+def deactivate_asset_dir(token):
+    _active_asset_dir.reset(token)
+
+
+def get_asset_dir(base_dir=None):
+    """The directory 'asset:' references are rooted at.
+
+    A library activated for this run wins outright - that is the server
+    telling the worker which workspace's assets this job uses. Otherwise
+    discovery mirrors the prompt library's - see workspace.discover_library
+    for the shared precedence (DW_ASSET_DIR, then a named workspace, then
+    ./assets, then a walk up from base_dir, then the workspace's assets/ as
+    the fallback).
+
+    Args:
+        base_dir: The workflow file's directory, when one anchors the search
+    """
+    active = _active_asset_dir.get()
+    if active:
+        return active
+
+    return discover_library(ASSETS_SUBDIR, ASSET_DIR_ENV_VAR, base_dir)
+
+
+def is_asset_reference(value):
+    """Whether a value references a file in the asset library."""
+    return isinstance(value, str) and value.startswith(ASSET_PREFIX)
+
+
+def asset_search_path(asset_dir=None, base_dir=None):
+    """Every directory an 'asset:' reference is looked for in, in order.
+
+    The workspace's own library first, then the read-only ones an entry
+    point put on the path (workspace.library_fallbacks - the assets a
+    --examples-dir tree brings with it), so an example workflow reaches the
+    media it ships with while an upload still lands in the workspace.
+
+    Args:
+        asset_dir: The first directory; defaults to get_asset_dir()
+        base_dir: The workflow file's directory, anchoring discovery when no
+            asset directory is configured
+    """
+    primary = asset_dir or get_asset_dir(base_dir)
+    return [primary] + library_fallbacks(ASSETS_SUBDIR, primary)
+
+
+def resolve_asset_reference(reference, asset_dir=None, base_dir=None):
+    """Resolve an 'asset:' reference to the file it names.
+
+    Args:
+        reference: The 'asset:name.ext' or 'asset:folder/name.ext' string
+        asset_dir: Directory the name is rooted at; defaults to get_asset_dir()
+        base_dir: The workflow file's directory, anchoring discovery when no
+            asset directory is configured
+
+    Returns:
+        The validated absolute path of the asset file
+
+    Raises:
+        InvalidInputError: If the name is not a valid asset name
+        PathTraversalError: If the name escapes the asset directory
+        ValueError: If no file exists under that name in any directory on
+            the search path
+    """
+    name = validate_asset_reference(reference.removeprefix(ASSET_PREFIX).strip())
+    roots = asset_search_path(asset_dir, base_dir)
+    for root in roots:
+        # Confined to the library it was found in: the name is joined onto a
+        # directory, so the containment check is what makes a name a name
+        # rather than a path
+        path = validate_path(os.path.join(root, name), root)
+        if os.path.isfile(path):
+            logger.debug(f"Resolved {reference} to {path}")
+            return path
+    searched = ", ".join(roots)
+    raise ValueError(
+        f"Asset '{name}' not found in {searched} - an 'asset:' reference "
+        f"names a file in the asset library, with its extension, like "
+        f"'asset:iris.jpg' or 'asset:gyre/frame_1.jpg'"
+    )
+
+
+def fetch_asset(reference, asset_dir=None, base_dir=None):
+    """The path an 'asset:' reference names, for whatever loads paths."""
+    return resolve_asset_reference(reference, asset_dir, base_dir)

dw/cache_blocks.json

@@ -0,0 +1,16 @@
+{
+    "$comment": "Transformer block metadata for diffusers' cache hooks. first_block, mag and layer_skip look their block class up in diffusers.hooks._helpers.TransformerBlockRegistry and raise if it is missing. Blocks listed here are registered on demand so those cache types work on models diffusers has not registered upstream. Fields mirror diffusers.hooks._helpers.TransformerBlockMetadata, plus dw's own encoder_hidden_states_argument_name. Verify against the block's forward before adding an entry: return_hidden_states_index is the position of hidden_states in the returned tuple, and is ignored when the block returns a bare tensor; return_encoder_hidden_states_index is null when the block does not return a second stream; hidden_states_argument_name must name the forward parameter carrying hidden_states; encoder_hidden_states_argument_name names the forward parameter the second returned stream comes from, and is only needed when that is not literally 'encoder_hidden_states' - see dw/cache_blocks.py.",
+    "blocks": {
+        "diffusers.models.transformers.transformer_minimax_h3.MiniMaxH3TransformerBlock": {
+            "return_hidden_states_index": 0,
+            "return_encoder_hidden_states_index": null,
+            "hidden_states_argument_name": "hidden_states"
+        },
+        "diffusers.models.transformers.transformer_ltx2.LTX2VideoTransformerBlock": {
+            "return_hidden_states_index": 0,
+            "return_encoder_hidden_states_index": 1,
+            "hidden_states_argument_name": "hidden_states",
+            "encoder_hidden_states_argument_name": "audio_hidden_states"
+        }
+    }
+}

dw/cache_blocks.py

@@ -0,0 +1,146 @@
+"""Register transformer block metadata diffusers' cache hooks are missing.
+
+The first_block, mag and layer_skip cache types look a model's transformer block
+class up in diffusers.hooks._helpers.TransformerBlockRegistry and raise
+ValueError when it is absent, so those caches are unavailable on any model
+diffusers has not registered upstream - even though the model itself supports
+enable_cache(). The blocks in cache_blocks.json fill that gap; entries become
+redundant, not wrong, once diffusers registers the same class itself.
+"""
+
+import json
+import logging
+from pathlib import Path
+
+from .type_helpers import load_type_from_name
+
+logger = logging.getLogger("dw")
+
+
+_REGISTRY_PATH = Path(__file__).parent / "cache_blocks.json"
+
+# Registration walks every entry, so do it once rather than per pipeline load
+_registered = False
+
+# The registry field that is dw's own rather than one of TransformerBlockMetadata's -
+# see _dual_stream_metadata_class for what it is for
+ARGUMENT_REMAP_KEY = "encoder_hidden_states_argument_name"
+
+# One subclass per metadata class, built on first use - the class it derives from
+# only exists once diffusers has been imported
+_dual_stream_classes = {}
+
+
+def _load_registry():
+    """Load the block metadata registry from the JSON file."""
+    with open(_REGISTRY_PATH) as f:
+        return json.load(f).get("blocks", {})
+
+
+def _dual_stream_metadata_class(metadata_class):
+    """A metadata class that can name the block's second stream itself.
+
+    A block that returns two streams has the second one read back out of its
+    forward arguments when the cache skips it, and diffusers looks that argument
+    up under the fixed name 'encoder_hidden_states' - the only two-stream shape it
+    registers upstream is text beside image. LTX-2's blocks return video beside
+    *audio* while also taking an 'encoder_hidden_states' of their own (the text
+    conditioning), so the fixed name silently reads the wrong tensor and feeds the
+    text embeddings back as the audio stream. Remapping the identifier is what
+    makes first_block caching correct on those blocks rather than merely quiet.
+
+    Args:
+        metadata_class: diffusers' TransformerBlockMetadata
+
+    Returns:
+        A subclass reading the second stream from the argument the entry names
+    """
+    if metadata_class not in _dual_stream_classes:
+
+        class DualStreamMetadata(metadata_class):
+            # Overridden per instance from the registry entry; the default keeps
+            # the class behaving exactly like the one it derives from
+            encoder_hidden_states_argument_name = "encoder_hidden_states"
+
+            def _get_parameter_from_args_kwargs(self, identifier, args=(), kwargs=None):
+                if identifier == "encoder_hidden_states":
+                    identifier = self.encoder_hidden_states_argument_name
+                return super()._get_parameter_from_args_kwargs(identifier, args, kwargs)
+
+        _dual_stream_classes[metadata_class] = DualStreamMetadata
+
+    return _dual_stream_classes[metadata_class]
+
+
+def build_metadata(metadata_class, entry):
+    """Build the metadata object one registry entry describes.
+
+    Args:
+        metadata_class: diffusers' TransformerBlockMetadata
+        entry: One block's fields from cache_blocks.json
+
+    Returns:
+        An instance of metadata_class, or of the dual-stream subclass when the
+        entry names the argument its second returned stream comes from
+    """
+    remapped = entry.get(ARGUMENT_REMAP_KEY)
+    fields = {k: v for k, v in entry.items() if k != ARGUMENT_REMAP_KEY}
+
+    if remapped is None:
+        return metadata_class(**fields)
+
+    metadata = _dual_stream_metadata_class(metadata_class)(**fields)
+    metadata.encoder_hidden_states_argument_name = remapped
+    return metadata
+
+
+def register_cache_blocks():
+    """Register any known-missing transformer blocks with diffusers.
+
+    Safe to call repeatedly and before any cache type - blocks diffusers already
+    knows are left alone, and a block whose class the installed diffusers does
+    not have is skipped rather than raising.
+    """
+    global _registered
+    if _registered:
+        return
+
+    try:
+        from diffusers.hooks._helpers import (
+            TransformerBlockMetadata,
+            TransformerBlockRegistry,
+        )
+    except ImportError:
+        # A diffusers without the cache hooks has nothing to register against
+        logger.debug(
+            "diffusers cache hook helpers unavailable, skipping block registration"
+        )
+        _registered = True
+        return
+
+    for class_name, metadata in _load_registry().items():
+        try:
+            # A server-owned registry name, registered, never constructed
+            block_class = load_type_from_name(class_name, constructed=False)
+        except (ImportError, AttributeError):
+            # The installed diffusers predates this model - its blocks are not
+            # missing from the registry, they do not exist
+            logger.debug(f"Cache block '{class_name}' not in this diffusers, skipping")
+            continue
+
+        # TransformerBlockRegistry.get raises rather than returning None, and an
+        # upstream registration is the more authoritative one either way
+        try:
+            TransformerBlockRegistry.get(block_class)
+            logger.debug(f"Cache block '{class_name}' already registered, leaving it")
+            continue
+        except ValueError:
+            pass
+
+        logger.debug(f"Registering cache block metadata for '{class_name}'")
+        TransformerBlockRegistry.register(
+            model_class=block_class,
+            metadata=build_metadata(TransformerBlockMetadata, metadata),
+        )
+
+    _registered = True