diffusers-workflow 0.4.0__py3-none-any.whl

dw/variables.py

@@ -0,0 +1,443 @@
+import copy
+import logging
+from .arguments import (
+    FROM_ARGUMENTS_KEY,
+    FROM_FILE_KEY,
+    FROM_PREVIOUS_RESULT_KEY,
+)
+from .security import (
+    validate_variable_name,
+    validate_string_input,
+    SecurityError,
+    MAX_VARIABLE_VALUE_LENGTH,
+)
+
+logger = logging.getLogger("dw")
+
+# The keys an object description names its media with. A "variable:" under one
+# of these that resolves to null has to stay present-and-null for
+# realize_object to read it as an omitted optional reference
+MEDIA_SOURCE_KEYS = (FROM_FILE_KEY, FROM_PREVIOUS_RESULT_KEY, FROM_ARGUMENTS_KEY)
+
+
+class VariableNotFoundError(ValueError):
+    """Raised when a workflow references a "variable:name" that isn't declared."""
+
+
+class VariableCycleError(ValueError):
+    """Raised when a list- or dict-valued variable references itself, directly
+    or through others, inside `resolve_variable_values`."""
+
+
+def _resolve_variable_reference(value, variables):
+    """
+    If value is a "variable:name" reference, look it up and return (True, resolved).
+    Otherwise return (False, None) so the caller knows to recurse instead.
+
+    Raises:
+        VariableNotFoundError: if the referenced name isn't in variables, naming the
+            variables that are actually available.
+    """
+    if isinstance(value, str) and value.startswith("variable:"):
+        variable_name = value.removeprefix("variable:")
+        logger.debug(f"Replacing variable reference: {variable_name}")
+        if variable_name not in variables:
+            available = ", ".join(sorted(variables.keys())) or "<none>"
+            message = f"Variable <{variable_name}> not found; available variables: {available}"
+            raise VariableNotFoundError(message)
+        return True, variables[variable_name]
+    return False, None
+
+
+def replace_variables(data, variables):
+    """
+    Recursively replaces variable references in data structures with their actual values.
+
+    Does not mutate its input - a new structure is returned and `data` is left as it
+    was passed in, so callers don't have to deep-copy defensively before calling.
+
+    Args:
+        data: The data structure (dict or list) containing variable references
+        variables: Dictionary of variable names and their values
+    Returns:
+        A new structure with "variable:name" references replaced. Any part of `data`
+        that isn't a dict/list/reference string is returned as-is.
+    """
+    if variables is None:
+        return data
+
+    logger.debug(f"Processing variables: {list(variables.keys())}")
+
+    # Handle lists - replace any "variable:name" strings with their values
+    if isinstance(data, list):
+        logger.debug(f"Processing list of length {len(data)}")
+        result = []
+        for item in data:
+            matched, resolved = _resolve_variable_reference(item, variables)
+            if matched:
+                result.append(resolved)
+            else:
+                # Recursively process nested structures
+                result.append(replace_variables(item, variables))
+        return result
+
+    # Handle dictionaries - replace values that are variable references
+    if isinstance(data, dict):
+        logger.debug(f"Processing dictionary with keys: {list(data.keys())}")
+        result = {}
+        for k, v in data.items():
+            matched, resolved = _resolve_variable_reference(v, variables)
+            if matched:
+                # A variable: reference that resolves to null is dropped as
+                # if the key had never been passed - an optional argument
+                # left at its unset default (e.g. select's threshold/index)
+                # shouldn't count as "present" just because a variable named
+                # it. A literal null written inline in the workflow JSON
+                # never goes through this branch, so its meaning elsewhere
+                # (a declared-but-unset default) is untouched.
+                #
+                # Except where the key itself is how an object description
+                # names its media: realize_object reads a present-and-null
+                # 'from_file'/'from_previous_result'/'from_arguments' as
+                # OMITTED - an optional reference this run was given nothing
+                # for (_names_no_media in arguments.py). Dropping the key
+                # there turns that into a media-less stub that reaches the
+                # pipeline instead.
+                if resolved is None and k not in MEDIA_SOURCE_KEYS:
+                    continue
+                result[k] = resolved
+            else:
+                # Recursively process nested structures in dictionary values
+                result[k] = replace_variables(v, variables)
+        return result
+
+    # Scalars (and anything else) pass through unchanged. copy.deepcopy guards
+    # against a caller mutating a returned mutable leaf (e.g. a PIL.Image or a
+    # list-typed variable's value) and having that reach back into `variables`.
+    return copy.deepcopy(data)
+
+
+def resolve_variable_values(variables):
+    """A copy of `variables` in which every "variable:name" inside a list-
+    or dict-valued variable is replaced by that variable's value.
+
+    A list-driven step reads its entries from a variable, and an entry that
+    says "from_file": "variable:character_a_voice" is how one variable sets
+    a voice in every shot the character speaks in. replace_variables only
+    walks the definition, so those references would reach the step as the
+    literal strings; this resolves them once, before realize_args, so a
+    reference type inside an entry is a type name by the time it is loaded.
+
+    Only list and dict values are walked. A scalar value that begins with
+    "variable:" is passed through as it always was.
+
+    Raises:
+        VariableNotFoundError: a reference names nothing declared
+        VariableCycleError: a value references itself, directly or through others
+    """
+    resolved = {}
+
+    def resolve(name, chain):
+        if name in resolved:
+            return resolved[name]
+        if name in chain:
+            loop = " -> ".join(chain[chain.index(name) :] + [name])
+            raise VariableCycleError(
+                f"Variable '{name}' references itself through: {loop}"
+            )
+        value = variables[name]
+        if isinstance(value, (list, dict)):
+            value = walk(value, chain + [name])
+        else:
+            value = copy.deepcopy(value)
+        resolved[name] = value
+        return value
+
+    def walk(node, chain):
+        matched, _ = _resolve_variable_reference(node, variables)
+        if matched:
+            return resolve(node.removeprefix("variable:"), chain)
+        if isinstance(node, list):
+            return [walk(item, chain) for item in node]
+        if isinstance(node, dict):
+            return {key: walk(item, chain) for key, item in node.items()}
+        return copy.deepcopy(node)
+
+    for name in variables:
+        resolve(name, [])
+    return resolved
+
+
+def undeclared_variable_references(definition):
+    """The "variable:name" references in a workflow definition that name no
+    entry of its `variables` - the ones `replace_variables` will refuse at run
+    time, found before anything loads.
+
+    Walks everything but `variables` itself, plus the inside of every list-
+    or dict-valued variable, the way `resolve_variable_values` and
+    `replace_variables` together do. Returns a list of (path, name) pairs,
+    path being where the reference sits (`steps[0].pipeline.arguments.prompt`)
+    and name what it asked for - which is the whole remainder of the string,
+    since a reference is the entire value and nothing is interpolated
+    around it.
+    """
+    declared = definition.get("variables") or {}
+    found = []
+
+    def walk(node, path):
+        if isinstance(node, str) and node.startswith("variable:"):
+            name = node.removeprefix("variable:")
+            if name not in declared:
+                found.append((path, name))
+        elif isinstance(node, dict):
+            for k, v in node.items():
+                walk(v, f"{path}.{k}" if path else k)
+        elif isinstance(node, list):
+            for i, v in enumerate(node):
+                walk(v, f"{path}[{i}]")
+
+    for key, value in definition.items():
+        if key != "variables":
+            walk(value, key)
+    if isinstance(declared, dict):
+        for name, value in declared.items():
+            if isinstance(value, (list, dict)):
+                walk(value, f"variables.{name}")
+    return found
+
+
+def _validated_strings(value):
+    """A copy of a list- or dict-valued caller argument with every string
+    leaf passed through `validate_string_input` - the same check a
+    top-level string argument gets in `set_variables` below. A for_each
+    entry's `prompt` or `from_file` is exactly as reachable to an attacker
+    as a top-level variable, and `isinstance(v, str)` alone would skip it.
+
+    Non-string leaves (numbers, bools, None, nested lists/dicts) pass
+    through unchanged; only str is validated.
+    """
+    if isinstance(value, str):
+        return validate_string_input(
+            value, max_length=MAX_VARIABLE_VALUE_LENGTH, allow_empty=True
+        )
+    if isinstance(value, list):
+        return [_validated_strings(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _validated_strings(item) for key, item in value.items()}
+    return value
+
+
+def set_variables(values, variables):
+    """
+    Sets the values of variables from a dictionary of new values with validation
+    Args:
+        values: Dictionary of new values to set
+        variables: Dictionary of existing variables with their default values/types
+    """
+    logger.debug(f"Setting variables: {list(values.keys())}")
+
+    if not isinstance(values, dict) or not isinstance(variables, dict):
+        logger.error("Both values and variables must be dictionaries")
+        raise TypeError("Both values and variables must be dictionaries")
+
+    for k, v in values.items():
+        try:
+            # Validate variable name
+            validated_name = validate_variable_name(k)
+
+            # The workflow must have already declared this variable (with a default
+            # value/type) - reject unknown names instead of raising a bare KeyError
+            if validated_name not in variables:
+                declared = ", ".join(sorted(variables.keys()))
+                logger.error(
+                    f"Unknown variable '{validated_name}'; declared variables: {declared}"
+                )
+                raise ValueError(
+                    f"Unknown variable '{validated_name}'; declared variables: {declared}"
+                )
+
+            # Validate string values - including strings nested inside a
+            # list- or dict-valued argument, which a for_each entry's
+            # prompt or from_file always is
+            if isinstance(v, str):
+                validated_value = validate_string_input(
+                    v, max_length=MAX_VARIABLE_VALUE_LENGTH, allow_empty=True
+                )
+            elif isinstance(v, (list, dict)):
+                validated_value = _validated_strings(v)
+            else:
+                validated_value = v
+
+            logger.debug(
+                f"Setting variable {validated_name} to value: {validated_value}"
+            )
+            # Use the type of the existing variable to convert the new value
+            variables[validated_name] = get_value(
+                validated_value, type(variables[validated_name]), validated_name
+            )
+
+        except SecurityError as e:
+            logger.error(f"Security validation failed for variable {k}: {e}")
+            raise
+
+
+def get_value(v, desired_type, name=None):
+    """
+    Converts a value to the desired type, with special handling for booleans
+    Args:
+        v: Value to convert
+        desired_type: Target type for conversion
+        name: Name of the variable being converted, used for error messages
+    Returns:
+        Converted value
+
+    Raises:
+        ValueError: if v cannot be converted to desired_type, naming the variable,
+            its target type, and the offending value.
+    """
+    logger.debug(f"Converting value {v} to type {desired_type}")
+
+    # A variable declared null is an optional one the workflow states no type
+    # for - passing a value to it is the expected case, not a suspicious one
+    if desired_type is None or desired_type is type(None):
+        logger.debug("Variable has no declared type, using the value as given")
+        return v
+
+    # A value already realized by an earlier step (an AudioTrack, an
+    # AudioVideo, a PIL.Image, ...) is a live object, not a JSON literal -
+    # 'previous_result:' resolves it to this before set_variables ever sees
+    # it, so a sibling step's task argument receives it unchanged. A
+    # sub-workflow's declared variable must too, rather than being coerced
+    # through the type of the variable's own default (usually a string
+    # 'asset:'/'output:' reference): desired_type(v) on one of these called
+    # str() on an AudioTrack and got its Python repr, which then reached
+    # slice_audio looking like a bogus path (#404)
+    if not isinstance(v, (str, int, float, bool, list, dict)):
+        logger.debug(
+            f"Value for variable '{name}' is an already-realized {type(v).__name__}; "
+            "using it as given"
+        )
+        return v
+
+    # Special handling for boolean string values - bool("0") and bool("no") are
+    # both truthy in Python, which would silently invert the user's intent, so
+    # only a known set of true/false spellings is accepted here
+    if isinstance(v, str) and desired_type is bool:
+        lowered = v.lower()
+        if lowered in ("true", "1", "yes", "on"):
+            return True
+        if lowered in ("false", "0", "no", "off"):
+            return False
+        var_label = name if name is not None else "<unknown>"
+        message = f"Cannot interpret '{v}' as true/false for variable '{var_label}'"
+        logger.error(message)
+        raise ValueError(message)
+
+    # Special handling for list string values - list("cat") would mangle the
+    # string into ['c', 'a', 't'], so a comma-separated string is split instead
+    if isinstance(v, str) and desired_type is list:
+        return [item.strip() for item in v.split(",")]
+
+    # A variable typed int by its default (e.g. `"score_gain": 1`) silently
+    # truncates a fractional override - int(0.3) == 0, with no error - which
+    # reads as a valid, if small, argument rather than the wrong type. Refuse
+    # it instead of truncating; the fix is to declare the default as a float
+    # (`1.0`) if the variable is meant to accept fractions
+    if desired_type is int and isinstance(v, float) and not v.is_integer():
+        var_label = name if name is not None else "<unknown>"
+        message = (
+            f"{var_label} {v!r} would be realized as {int(v)}: this variable "
+            f"is typed integer by its default; declare its default as a float "
+            f"(e.g. {float(int(v))!r}) to accept fractional values"
+        )
+        logger.error(message)
+        raise ValueError(message)
+
+    # A string cannot be coerced into a dict or a None - dict('/a/b.png') is
+    # nonsense, NoneType('x') a TypeError. Those defaults are how media
+    # variables ({'location': ...}) and optional inputs (null) are declared,
+    # and a string override is a path or a reference that realize_args
+    # resolves later, so it passes through as written
+    if isinstance(v, str) and desired_type in (dict, type(None)):
+        return v
+
+    # A dict or list cannot be coerced into a string - str({...}) never
+    # raises, it just stringifies the Python repr, so a caller passing the
+    # old {"location": ...} object for a variable now declared as a plain
+    # string (e.g. templates/ltx2/keyframes' first_image/last_image, #433)
+    # sailed through validation and only failed at run time with a path
+    # built from the dict's repr
+    if isinstance(v, (dict, list)) and desired_type is str:
+        var_label = name if name is not None else "<unknown>"
+        message = (
+            f"Cannot convert variable '{var_label}' value {v!r} to type str: "
+            f"this variable takes a plain string (a URL, or an 'asset:'/"
+            f"'output:' reference), not an object"
+        )
+        logger.error(message)
+        raise ValueError(message)
+
+    # Attempt type conversion. A failure here is surfaced immediately with a clear,
+    # named error instead of silently passing the unconverted value through - letting
+    # it through would fail several layers later inside diffusers/torch with a
+    # confusing traceback that doesn't mention the variable at fault.
+    try:
+        converted = desired_type(v)
+        logger.debug(f"Successfully converted to {desired_type.__name__}: {converted}")
+        return converted
+    except Exception as e:
+        var_label = name if name is not None else "<unknown>"
+        message = (
+            f"Cannot convert variable '{var_label}' value {v!r} to type "
+            f"{desired_type.__name__}: {e}"
+        )
+        logger.error(message)
+        raise ValueError(message) from e
+
+
+def argument_errors(definition, arguments):
+    """What is wrong with a caller's `arguments` for this workflow, before
+    anything is queued.
+
+    Exactly the check `set_variables` makes at the top of a run - an
+    undeclared name, a value that will not coerce to the type the default
+    declares - made against a copy, so nothing is mutated and the answer
+    costs no GPU time. A workflow that declares no variables at all takes no
+    arguments: today those are dropped in silence (`Workflow.run` only
+    substitutes when a `variables` block exists), which is the one case the
+    run itself does not report.
+
+    Args:
+        definition: A workflow definition
+        arguments: The caller's argument dict; empty or None means no errors
+
+    Returns:
+        [{"path": "arguments.<name>", "message": str}, ...], one per bad
+        argument, in the order they were given
+    """
+    if not arguments:
+        return []
+    if not isinstance(arguments, dict):
+        return [{"path": "arguments", "message": "arguments must be an object"}]
+
+    declared = definition.get("variables") if isinstance(definition, dict) else None
+    if not isinstance(declared, dict) or not declared:
+        return [
+            {
+                "path": f"arguments.{name}",
+                "message": f"This workflow declares no variables, so '{name}' "
+                "has nowhere to land - it would be ignored by the run",
+            }
+            for name in arguments
+        ]
+
+    errors = []
+    for name, value in arguments.items():
+        # One at a time against a fresh copy, so each bad argument is
+        # reported with its own name rather than the first one stopping the
+        # rest from being checked
+        try:
+            set_variables({name: value}, copy.deepcopy(declared))
+        except (ValueError, TypeError, SecurityError) as e:
+            errors.append({"path": f"arguments.{name}", "message": str(e)})
+    return errors

dw/video_extensions.py

@@ -0,0 +1,141 @@
+"""The extension of a `video` argument, checked before the run when it can be.
+
+`fetch_video` (`dw/arguments.py`) loads every argument named `video` or
+`*_video`, and every `{"media_type": "video", "location": ...}` reference
+whatever it is named, as a video file - and refuses one whose extension is
+not `ALLOWED_VIDEO_EXTENSIONS` there, at run time. A still image handed to
+one of those arguments (`"video": "asset:sheet.png"`) validated clean and
+then failed the job in seconds with "Video file extension not allowed:
+.png", after `get_task("loop_frames")` had described the same argument as
+taking "a still image, or frames in any shape a result carries" (#347).
+
+Moved here, into `validation_errors`, for exactly the cases whose extension
+is knowable without running anything: a literal file path, or an
+`asset:`/`output:` reference whose name carries one - `resolve_path_references`
+turns either into a real path before `fetch_video` ever sees it, and the
+extension survives that unchanged. A `previous_result:` (or any reference
+`expand_for_each` left unresolved) names no file yet and is left to the
+existing run-time check, and a URL is left alone too - `fetch_video` never
+gates a URL's extension, so refusing one here would refuse something the run
+itself accepts.
+"""
+
+from .arguments import (
+    CONSTANT_PREFIX,
+    PROMPT_PREFIX,
+    is_media_reference,
+)
+from .for_each import MEMBER_SEPARATOR, render_path
+from .security import ALLOWED_IMAGE_EXTENSIONS, ALLOWED_VIDEO_EXTENSIONS
+
+# Left to the run-time check: not yet resolved to anything an extension can
+# be read off, at the point validation walks the expanded definition
+_UNRESOLVED_PREFIXES = ("previous_result:", "variable:", "item:", "gather:")
+
+
+def _is_video_key(key):
+    return isinstance(key, str) and (key == "video" or key.endswith("_video"))
+
+
+def _extension_problem(value):
+    """Why this string's extension is not one `fetch_video` will accept, or
+    None - including None for anything whose extension is not yet knowable."""
+    if not isinstance(value, str):
+        return None
+    if value.startswith(("http://", "https://")):
+        return None
+    if value.startswith(_UNRESOLVED_PREFIXES):
+        return None
+    if value.startswith(CONSTANT_PREFIX) or value.startswith(PROMPT_PREFIX):
+        return None
+    ext = value.rsplit(".", 1)
+    if len(ext) != 2 or not ext[1] or "/" in ext[1]:
+        return None
+    ext = f".{ext[1].lower()}"
+    if ext in ALLOWED_VIDEO_EXTENSIONS:
+        return None
+    if ext in ALLOWED_IMAGE_EXTENSIONS:
+        return (
+            f"'{value}' is a still image, and a video argument loads video "
+            f"files - pass it as "
+            f'{{"media_type": "image", "location": "{value}"}} to load it as '
+            f"a still, or reference a prior image step with previous_result:"
+        )
+    return f"Video file extension not allowed: {ext}"
+
+
+def _video_key_values(value, path):
+    """Strings a video-named key hands to `fetch_video` - itself, or each
+    element of a list of them. A dict here is a `media_type` reference and
+    is walked by `_video_values` instead, same as `fetch_video` handles it."""
+    if isinstance(value, str):
+        yield path, value
+    elif isinstance(value, list):
+        for index, item in enumerate(value):
+            yield from _video_key_values(item, path + (index,))
+    elif isinstance(value, dict):
+        yield from _video_values(value, path)
+
+
+def _video_values(value, path):
+    """Every string this walk can attribute to a location `fetch_video`
+    will load, paired with the path it sits at."""
+    if isinstance(value, dict):
+        if is_media_reference(value) and value.get("media_type") == "video":
+            location = value.get("location")
+            if isinstance(location, str):
+                yield path + ("location",), location
+            return
+        for key, item in value.items():
+            if _is_video_key(key):
+                yield from _video_key_values(item, path + (key,))
+            else:
+                yield from _video_values(item, path + (key,))
+    elif isinstance(value, list):
+        for index, item in enumerate(value):
+            yield from _video_values(item, path + (index,))
+
+
+def video_extension_errors(workflow_definition, source_indices=None):
+    """Every video argument whose extension `fetch_video` will refuse, for
+    every case that extension is knowable before the run, as
+    [{path, message}].
+
+    Walks the substituted, expanded definition, the same convention
+    `reference_name_errors` and `task_argument_errors` follow: `source_indices`
+    maps an expanded step back to the one the author wrote, and a path inside
+    a `for_each` member names the member.
+    """
+    steps = workflow_definition.get("steps")
+    if not isinstance(steps, list):
+        return []
+
+    errors = []
+    for index, step in enumerate(steps):
+        if not isinstance(step, dict):
+            continue
+        source = (
+            source_indices[index]
+            if source_indices is not None and index < len(source_indices)
+            else index
+        )
+        name = step.get("name")
+        where = (
+            f" in member '{name}'"
+            if isinstance(name, str) and MEMBER_SEPARATOR in name
+            else ""
+        )
+        for path, value in _video_values(step, ()):
+            problem = _extension_problem(value)
+            if problem is None:
+                continue
+            errors.append(
+                {
+                    "path": render_path(("steps", source) + path),
+                    "message": f"{problem}{where}",
+                }
+            )
+    return errors
+
+
+__all__ = ["video_extension_errors"]