diffusers-workflow 0.4.0a3__py3-none-any.whl

dw/pipeline_processors/remote.py

@@ -0,0 +1,18 @@
+import torch
+from huggingface_hub import get_token
+import requests
+import io
+
+
+def remote_text_encoder(prompts, url, device):
+    response = requests.post(
+        url,
+        json={"prompt": prompts},
+        headers={
+            "Authorization": f"Bearer {get_token()}",
+            "Content-Type": "application/json",
+        },
+    )
+    prompt_embeds = torch.load(io.BytesIO(response.content))
+
+    return prompt_embeds.to(device)

dw/previous_results.py

@@ -0,0 +1,259 @@
+import logging
+from itertools import product
+
+from .arguments import (
+    FROM_PREVIOUS_RESULT_KEY,
+    PREVIOUS_RESULT_PREFIX,
+    build_objects,
+)
+
+logger = logging.getLogger("dw")
+
+# Maximum number of iterations to prevent resource exhaustion
+MAX_ITERATIONS = 10000
+
+
+def get_iterations(argument_template, previous_results):
+    """Generate argument combinations using previous task results.
+
+    Takes a template of arguments and expands any references to previous results
+    into all possible combinations of those results.
+
+    Args:
+        argument_template: Dict or list containing argument definitions
+        previous_results: Dict of results from previously executed steps
+
+    Returns:
+        List of argument dictionaries, one for each possible combination
+    """
+    # Special case: if template is a list, use it directly without processing
+    if isinstance(argument_template, list):
+        logger.debug("Using list argument template directly")
+        return argument_template
+
+    # Find any references to previous results in the template
+    # Returns dict of {arg_key: result_reference}
+    result_refs = find_previous_result_refs(argument_template)
+
+    # If no references found, return the template as-is
+    if not result_refs:
+        logger.debug("No result references found in template")
+        # Shallow copy: realize_args may have already loaded large media
+        # (PIL images, full video frame lists) into the template, so a deep
+        # copy would multiply memory use. Contract: iteration dicts may only
+        # be mutated at the top level (key pop/assign); nested values are
+        # shared across iterations and must never be mutated in place.
+        return [dict(argument_template)]
+
+    logger.debug(f"Found {len(result_refs)} result references: {result_refs}")
+
+    # Create a dictionary mapping each reference path to its possible values
+    # Example: {('image',): [img1, img2], ('prompt',): ['text1', 'text2']}
+    ref_results = {
+        ref_path: list(get_previous_results(previous_results, ref_value))
+        for ref_path, ref_value in result_refs.items()
+    }
+
+    # Generate all possible combinations of argument values
+    keys = list(ref_results.keys())
+    iterations = []
+
+    # Use itertools.product to create cartesian product of all possible values
+    # Example: if ref_results has 2 images and 2 prompts, creates 4 combinations
+    for values in product(*[ref_results[k] for k in keys]):
+        # Create fresh shallow copy of template for each combination.
+        # Nested values (e.g. loaded PIL images, video frame lists) are
+        # shared across iterations, not deep-copied, to avoid multiplying
+        # media memory usage by the iteration count. Contract: iteration
+        # dicts may only be mutated at the top level (key pop/assign);
+        # nested values must never be mutated in place.
+        arguments = dict(argument_template)
+
+        # Replace each reference with its actual value
+        for path, value in zip(keys, values):
+            # Handle nested dictionary properties
+            # If value is dict and contains the key we're looking for, use that property
+            key = path[-1]
+            arguments = substitute_at_path(
+                arguments,
+                path,
+                value[key] if isinstance(value, dict) and key in value else value,
+            )
+
+        # Now that the media exists, build the objects that were waiting for it -
+        # a reference constructed from a step's output rather than from a file
+        iterations.append(build_objects(arguments))
+
+    # Safety check to prevent cartesian product explosion
+    if len(iterations) > MAX_ITERATIONS:
+        raise ValueError(
+            f"Too many iterations generated: {len(iterations)} exceeds maximum of {MAX_ITERATIONS}. "
+            f"This usually indicates too many previous_result references creating a cartesian product. "
+            f"Consider reducing the number of multi-value results or splitting into multiple steps."
+        )
+
+    logger.debug(f"Generated {len(iterations)} argument combinations")
+    return iterations
+
+
+def get_previous_results(previous_results, previous_result_name):
+    """Retrieve results or specific properties from previous tasks.
+
+    Args:
+        previous_results: Dict of results from previous steps
+        previous_result_name: String identifying the result, optionally with property
+                            Format: "step_name" or "step_name.property_name"
+
+    Returns:
+        List of results or specific properties from the referenced step
+    """
+    # Step names are unrestricted strings and may themselves contain dots
+    # (e.g. "v1.0"), so resolve against the known step names rather than
+    # blindly splitting on the first/only ".".
+
+    # Exact match: the whole reference is a known step name, no property.
+    if previous_result_name in previous_results:
+        logger.debug(f"Getting all artifacts from result {previous_result_name}")
+        return previous_results[previous_result_name].get_artifacts()
+
+    if "." not in previous_result_name:
+        raise KeyError(
+            f"Previous result '{previous_result_name}' not found. Available results: {list(previous_results.keys())}"
+        )
+
+    # Find the longest known step name that is a prefix of the reference
+    # followed by ".", and treat the remainder as the property name.
+    result_name = max(
+        (
+            name
+            for name in previous_results
+            if previous_result_name.startswith(name + ".")
+        ),
+        key=len,
+        default=None,
+    )
+
+    if result_name is None:
+        raise KeyError(
+            f"Previous result '{previous_result_name}' not found. Available results: {list(previous_results.keys())}"
+        )
+
+    property_name = previous_result_name[len(result_name) + 1 :]
+    logger.debug(f"Getting property {property_name} from result {result_name}")
+    return previous_results[result_name].get_artifact_properties(property_name)
+
+
+def resolve_chain_prompts(step_action, previous_results):
+    """Resolve a pipeline chain's per-segment prompts against previous results.
+
+    A chain's "prompts" list is not part of the step's argument template, so the
+    cartesian pass that expands "previous_result:" everywhere else never reaches
+    it. That matters for a chain whose opening segment is written by a different
+    step from the ones that continue it - a continuation prompt declares a video
+    reference the first segment does not have.
+
+    Each entry resolves independently and yields one prompt, so this never
+    multiplies iterations the way an argument reference does; a reference that
+    produced several artifacts uses the first.
+
+    The resolved list is left on the step action for run_chain to pick up, and
+    nothing happens at all for a pipeline without chain prompts.
+    """
+    definition = getattr(step_action, "pipeline_definition", None)
+    if not isinstance(definition, dict):
+        return
+
+    chain = definition.get("chain", None) or {}
+    prompts = chain.get("prompts", None)
+    if not prompts:
+        return
+
+    resolved = []
+    for entry in prompts:
+        if isinstance(entry, str) and entry.startswith("previous_result:"):
+            artifacts = get_previous_results(
+                previous_results, entry.removeprefix("previous_result:")
+            )
+            if not artifacts:
+                raise ValueError(f"Chain prompt reference '{entry}' produced no result")
+            if len(artifacts) > 1:
+                logger.warning(
+                    f"Chain prompt reference '{entry}' produced {len(artifacts)} "
+                    f"results - using the first"
+                )
+            entry = artifacts[0]
+        resolved.append(entry)
+
+    step_action.chain_prompts = resolved
+
+
+def find_previous_result_refs(arguments):
+    """Find all values in an argument structure that reference previous results.
+
+    A reference is written either as a value with the "previous_result:" prefix, or
+    as the step name a 'from_previous_result' object description is built from. Both
+    are found at any depth: an argument that takes a constructed object holds it
+    inside a list - MiniMax-H3's 'references' - so the reference is nested rather
+    than sitting at the top of the arguments.
+
+    Args:
+        arguments: Dictionary of argument definitions
+
+    Returns:
+        Dict mapping the path of each reference to the result name it names. A path
+        is the tuple of keys and list indices that reaches the value, so a top-level
+        {'image': 'previous_result:step1'} comes back as {('image',): 'step1'}
+    """
+    found = {}
+    _collect_refs(arguments, (), found)
+    return found
+
+
+def _collect_refs(value, path, found):
+    """Walk an argument structure, collecting every reference by its path."""
+    if isinstance(value, dict):
+        for key, item in value.items():
+            # The object description names its step bare, the way it would name a
+            # file - the prefix would only repeat what the key already says
+            if key == FROM_PREVIOUS_RESULT_KEY and isinstance(item, str):
+                found[path + (key,)] = item
+            else:
+                _collect_refs(item, path + (key,), found)
+
+    elif isinstance(value, list):
+        for index, item in enumerate(value):
+            _collect_refs(item, path + (index,), found)
+
+    elif isinstance(value, str) and value.startswith(PREVIOUS_RESULT_PREFIX):
+        found[path] = value[len(PREVIOUS_RESULT_PREFIX) :]
+
+
+def substitute_at_path(container, path, value):
+    """A copy of container with value placed at path.
+
+    Only the containers along the path are copied. Everything beside them stays
+    shared, which is the same contract the top-level copy keeps: iterations share
+    their nested values, so a substitution deep in one must not be visible in the
+    others.
+
+    Args:
+        container: The dict or list to substitute into
+        path: Tuple of keys and indices reaching the value to replace
+        value: What to put there
+
+    Returns:
+        The copied container
+    """
+    key = path[0]
+    replacement = (
+        value if len(path) == 1 else substitute_at_path(container[key], path[1:], value)
+    )
+
+    if isinstance(container, list):
+        copied = list(container)
+        copied[key] = replacement
+        return copied
+
+    copied = dict(container)
+    copied[key] = replacement
+    return copied

dw/prompt_weighting.py

@@ -0,0 +1,378 @@
+"""
+Prompt weighting and long prompt support for diffusers pipelines.
+
+Parses A1111-style prompt syntax: (word:1.5) for emphasis, [word] for de-emphasis,
+((word)) for nested weighting. Supports prompts longer than the 77-token CLIP limit.
+
+Produces prompt_embeds tensors that replace the prompt string argument in pipeline calls.
+
+Based on sd_embed by Andrew Zhu (https://github.com/xhinker/sd_embed)
+License: Apache 2.0
+"""
+
+import re
+import gc
+import logging
+from typing import Tuple
+
+import torch
+from transformers import CLIPTokenizer, T5Tokenizer
+
+logger = logging.getLogger("dw")
+
+
+# ---------------------------------------------------------------------------
+# Prompt parser — A1111-style (word:weight) syntax
+# ---------------------------------------------------------------------------
+
+_re_attention = re.compile(
+    r"""
+\\\(|
+\\\)|
+\\\[|
+\\]|
+\\\\|
+\\|
+\(|
+\[|
+:\s*([+-]?[.\d]+)\s*\)|
+\)|
+]|
+[^\\()\[\]:]+|
+:
+""",
+    re.X,
+)
+
+_re_break = re.compile(r"\s*\bBREAK\b\s*", re.S)
+
+
+def parse_prompt_attention(text):
+    """Parse a prompt string with attention weights.
+
+    Syntax:
+        (abc)       — weight 1.1
+        (abc:1.5)   — weight 1.5
+        ((abc))     — weight 1.21 (1.1 * 1.1)
+        [abc]       — weight 1/1.1 ≈ 0.91
+        \\( \\)     — literal parens
+
+    Returns list of [text, weight] pairs.
+    """
+    res = []
+    round_brackets = []
+    square_brackets = []
+    round_bracket_multiplier = 1.1
+    square_bracket_multiplier = 1 / 1.1
+
+    def multiply_range(start_position, multiplier):
+        for p in range(start_position, len(res)):
+            res[p][1] *= multiplier
+
+    for m in _re_attention.finditer(text):
+        text_match = m.group(0)
+        weight = m.group(1)
+
+        if text_match.startswith("\\"):
+            res.append([text_match[1:], 1.0])
+        elif text_match == "(":
+            round_brackets.append(len(res))
+        elif text_match == "[":
+            square_brackets.append(len(res))
+        elif weight is not None and len(round_brackets) > 0:
+            multiply_range(round_brackets.pop(), float(weight))
+        elif text_match == ")" and len(round_brackets) > 0:
+            multiply_range(round_brackets.pop(), round_bracket_multiplier)
+        elif text_match == "]" and len(square_brackets) > 0:
+            multiply_range(square_brackets.pop(), square_bracket_multiplier)
+        else:
+            parts = re.split(_re_break, text_match)
+            for i, part in enumerate(parts):
+                if i > 0:
+                    res.append(["BREAK", -1])
+                res.append([part, 1.0])
+
+    for pos in round_brackets:
+        multiply_range(pos, round_bracket_multiplier)
+    for pos in square_brackets:
+        multiply_range(pos, square_bracket_multiplier)
+
+    if len(res) == 0:
+        res = [["", 1.0]]
+
+    # merge runs of identical weights
+    i = 0
+    while i + 1 < len(res):
+        if res[i][1] == res[i + 1][1]:
+            res[i][0] += res[i + 1][0]
+            res.pop(i + 1)
+        else:
+            i += 1
+
+    return res
+
+
+# ---------------------------------------------------------------------------
+# Tokenization helpers
+# ---------------------------------------------------------------------------
+
+
+def _tokenize_clip_with_weights(clip_tokenizer: CLIPTokenizer, prompt: str):
+    """Tokenize with CLIP and return (token_ids, weights)."""
+    if not prompt:
+        prompt = "empty"
+
+    texts_and_weights = parse_prompt_attention(prompt)
+    text_tokens, text_weights = [], []
+    for word, weight in texts_and_weights:
+        token = clip_tokenizer(word, truncation=False).input_ids[1:-1]
+        text_tokens.extend(token)
+        text_weights.extend([weight] * len(token))
+    return text_tokens, text_weights
+
+
+def _tokenize_t5_with_weights(t5_tokenizer: T5Tokenizer, prompt: str):
+    """Tokenize with T5 and return (token_ids, weights)."""
+    if not prompt:
+        prompt = "empty"
+
+    texts_and_weights = parse_prompt_attention(prompt)
+    text_tokens, text_weights = [], []
+    for word, weight in texts_and_weights:
+        token = t5_tokenizer(word, truncation=False, add_special_tokens=True).input_ids
+        text_tokens.extend(token)
+        text_weights.extend([weight] * len(token))
+    return text_tokens, text_weights
+
+
+def _group_tokens_and_weights(token_ids, weights, pad_last_block=True):
+    """Group tokens into 77-token chunks with BOS/EOS padding."""
+    bos, eos = 49406, 49407
+
+    new_token_ids = []
+    new_weights = []
+
+    # work on copies to avoid mutating originals
+    token_ids = list(token_ids)
+    weights = list(weights)
+
+    while len(token_ids) >= 75:
+        head_tokens = [token_ids.pop(0) for _ in range(75)]
+        head_weights = [weights.pop(0) for _ in range(75)]
+        new_token_ids.append([bos] + head_tokens + [eos])
+        new_weights.append([1.0] + head_weights + [1.0])
+
+    if len(token_ids) > 0:
+        padding_len = 75 - len(token_ids) if pad_last_block else 0
+        new_token_ids.append([bos] + token_ids + [eos] * padding_len + [eos])
+        new_weights.append([1.0] + weights + [1.0] * padding_len + [1.0])
+
+    return new_token_ids, new_weights
+
+
+# ---------------------------------------------------------------------------
+# Pipeline-specific weighted embedding functions
+# ---------------------------------------------------------------------------
+
+
+def _get_device(pipeline):
+    """Get the appropriate compute device for the pipeline."""
+    device = pipeline.device
+    # Offloaded pipelines report cpu (model offload) or meta (sequential offload)
+    if device is not None and device.type not in ("cpu", "meta"):
+        return device
+
+    # Fall back to the device dw is running on, which honors the DW_DEVICE and
+    # settings overrides - never a hardcoded accelerator
+    from . import get_device
+
+    return torch.device(get_device())
+
+
+def _hook_managed(module):
+    """Whether accelerate offload hooks own this module's placement.
+
+    Moving a hooked module manually fights the hooks - and raises outright on the
+    meta tensors sequential offload leaves behind. The hooks bring the module to
+    its execution device on forward, so no manual move is needed.
+    """
+    return getattr(module, "_hf_hook", None) is not None
+
+
+def get_weighted_text_embeddings_flux(
+    pipe,
+    prompt: str = "",
+    prompt2: str = None,
+    device=None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Generate weighted text embeddings for Flux pipelines.
+
+    Supports long prompts (beyond 77 tokens) and A1111-style weighting syntax.
+
+    Args:
+        pipe: A loaded FluxPipeline with tokenizer, tokenizer_2, text_encoder, text_encoder_2
+        prompt: Primary prompt with optional weighting syntax
+        prompt2: Optional second prompt for T5 encoder (defaults to prompt)
+        device: Target device override
+
+    Returns:
+        (prompt_embeds, pooled_prompt_embeds) — pass directly to pipe() as kwargs
+    """
+    prompt2 = prompt if prompt2 is None else prompt2
+
+    target_device = device if device is not None else _get_device(pipe)
+
+    # Move text encoders to device if the pipeline sits on the CPU - unless
+    # offload hooks manage them, in which case they move themselves on forward
+    encoders_moved = False
+    if pipe.device.type == "cpu" and not (
+        _hook_managed(pipe.text_encoder) or _hook_managed(pipe.text_encoder_2)
+    ):
+        pipe.text_encoder.to(target_device)
+        pipe.text_encoder_2.to(target_device)
+        encoders_moved = True
+
+    # Tokenize with CLIP (tokenizer 1) for pooled embeddings
+    prompt_tokens, prompt_weights = _tokenize_clip_with_weights(pipe.tokenizer, prompt)
+    prompt_token_groups, _ = _group_tokens_and_weights(prompt_tokens, prompt_weights)
+
+    # Generate pooled CLIP embeddings (mean across token groups)
+    pool_embeds_list = []
+    for token_group in prompt_token_groups:
+        token_tensor = torch.tensor(
+            [token_group], dtype=torch.long, device=target_device
+        )
+        with torch.no_grad():
+            embeds = pipe.text_encoder(token_tensor, output_hidden_states=False)
+        pool_embeds_list.append(embeds.pooler_output.squeeze(0))
+
+    pooled_prompt_embeds = torch.stack(pool_embeds_list, dim=0)
+    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=0, keepdim=True)
+    pooled_prompt_embeds = pooled_prompt_embeds.to(
+        dtype=pipe.text_encoder.dtype, device=target_device
+    )
+
+    # Tokenize with T5 (tokenizer 2) for main prompt embeddings
+    prompt_tokens_2, prompt_weights_2 = _tokenize_t5_with_weights(
+        pipe.tokenizer_2, prompt2
+    )
+
+    token_tensor_2 = torch.tensor([prompt_tokens_2], dtype=torch.long)
+    with torch.no_grad():
+        t5_embeds = pipe.text_encoder_2(token_tensor_2.to(target_device))[0].squeeze(0)
+    t5_embeds = t5_embeds.to(device=target_device)
+
+    # Apply per-token weights to T5 embeddings
+    for i in range(len(prompt_weights_2)):
+        if prompt_weights_2[i] != 1.0:
+            t5_embeds[i] = t5_embeds[i] * prompt_weights_2[i]
+
+    prompt_embeds = t5_embeds.unsqueeze(0).to(
+        dtype=pipe.text_encoder_2.dtype, device=target_device
+    )
+
+    # Release encoders back to CPU if we moved them
+    if encoders_moved:
+        pipe.text_encoder.to("cpu")
+        pipe.text_encoder_2.to("cpu")
+        gc.collect()
+
+        from . import empty_device_cache
+
+        empty_device_cache()
+
+    return prompt_embeds, pooled_prompt_embeds
+
+
+# ---------------------------------------------------------------------------
+# Dispatcher — selects the right function based on pipeline type
+# ---------------------------------------------------------------------------
+
+# Map pipeline class names to their embedding functions
+_PIPELINE_FUNCTIONS = {
+    "FluxPipeline": get_weighted_text_embeddings_flux,
+    "FluxImg2ImgPipeline": get_weighted_text_embeddings_flux,
+    "FluxInpaintPipeline": get_weighted_text_embeddings_flux,
+    "FluxControlNetPipeline": get_weighted_text_embeddings_flux,
+}
+
+# The encoder stack get_weighted_text_embeddings_flux drives - CLIP for pooled
+# embeddings, T5 for the main ones
+_FLUX_ENCODER_STACK = ("tokenizer", "tokenizer_2", "text_encoder", "text_encoder_2")
+
+
+def _select_embedding_function(pipeline):
+    """The weighted-embedding function for a pipeline, or None.
+
+    Exact class names first; any other Flux-family pipeline (FluxKontext,
+    FluxFill, a user subclass) carrying the same CLIP+T5 encoder stack uses the
+    flux function too, so support does not lag every new variant diffusers adds.
+    """
+    embed_fn = _PIPELINE_FUNCTIONS.get(pipeline.__class__.__name__)
+    if embed_fn is not None:
+        return embed_fn
+
+    if pipeline.__class__.__name__.startswith("Flux") and all(
+        getattr(pipeline, name, None) is not None for name in _FLUX_ENCODER_STACK
+    ):
+        logger.info(
+            f"{pipeline.__class__.__name__} carries the Flux encoder stack - "
+            f"applying flux prompt weighting"
+        )
+        return get_weighted_text_embeddings_flux
+
+    return None
+
+
+def apply_prompt_weighting(pipeline, arguments, device=None):
+    """Apply prompt weighting to pipeline arguments if the prompt contains weight syntax.
+
+    Checks if the prompt uses weighting syntax. If so, generates weighted embeddings
+    and replaces the prompt string with embedding tensors in the arguments dict.
+
+    Args:
+        pipeline: The loaded diffusers pipeline
+        arguments: Mutable dict of pipeline call arguments
+        device: The device the pipeline runs on - embeddings are created there.
+            Defaults to the pipeline's own device, or the dw device when offloading
+            parks the pipeline on the CPU
+
+    Returns:
+        True if weighting was applied, False if prompt was left as-is.
+    """
+    prompt = arguments.get("prompt", None)
+    if prompt is None or not isinstance(prompt, str):
+        return False
+
+    # Quick check: does the prompt contain any weighting syntax?
+    if "(" not in prompt and "[" not in prompt:
+        return False
+
+    class_name = pipeline.__class__.__name__
+    embed_fn = _select_embedding_function(pipeline)
+    if embed_fn is None:
+        logger.warning(
+            f"Prompt weighting not supported for {class_name}. "
+            f"Supported: {', '.join(_PIPELINE_FUNCTIONS.keys())} and Flux-family "
+            f"pipelines with the CLIP+T5 encoder stack. "
+            f"Passing prompt as plain text."
+        )
+        return False
+
+    logger.info(f"Applying prompt weighting for {class_name}")
+    prompt2 = arguments.pop("prompt_2", None)
+    prompt_str = arguments.pop("prompt")
+
+    prompt_embeds, pooled_prompt_embeds = embed_fn(
+        pipeline, prompt=prompt_str, prompt2=prompt2, device=device
+    )
+
+    arguments["prompt_embeds"] = prompt_embeds
+    arguments["pooled_prompt_embeds"] = pooled_prompt_embeds
+
+    # Remove negative_prompt if present — can't mix string and embeds
+    if "negative_prompt" in arguments:
+        logger.debug("Removing negative_prompt (incompatible with prompt_embeds)")
+        arguments.pop("negative_prompt")
+
+    return True