Pipeline states

mindone.diffusers.modular_pipelines.modular_pipeline.PipelineState dataclass

[PipelineState] stores the state of a pipeline. It is used to pass data between pipeline blocks.

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

@dataclass
class PipelineState:
    """
    [`PipelineState`] stores the state of a pipeline. It is used to pass data between pipeline blocks.
    """

    values: Dict[str, Any] = field(default_factory=dict)
    kwargs_mapping: Dict[str, List[str]] = field(default_factory=dict)

    def set(self, key: str, value: Any, kwargs_type: str = None):
        """
        Add a value to the pipeline state.

        Args:
            key (str): The key for the value
            value (Any): The value to store
            kwargs_type (str): The kwargs_type with which the value is associated
        """
        self.values[key] = value

        if kwargs_type is not None:
            if kwargs_type not in self.kwargs_mapping:
                self.kwargs_mapping[kwargs_type] = [key]
            else:
                self.kwargs_mapping[kwargs_type].append(key)

    def get(self, keys: Union[str, List[str]], default: Any = None) -> Union[Any, Dict[str, Any]]:
        """
        Get one or multiple values from the pipeline state.

        Args:
            keys (Union[str, List[str]]): Key or list of keys for the values
            default (Any): The default value to return if not found

        Returns:
            Union[Any, Dict[str, Any]]: Single value if keys is str, dictionary of values if keys is list
        """
        if isinstance(keys, str):
            return self.values.get(keys, default)
        return {key: self.values.get(key, default) for key in keys}

    def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]:
        """
        Get all values with matching kwargs_type.

        Args:
            kwargs_type (str): The kwargs_type to filter by

        Returns:
            Dict[str, Any]: Dictionary of values with matching kwargs_type
        """
        value_names = self.kwargs_mapping.get(kwargs_type, [])
        return self.get(value_names)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert PipelineState to a dictionary.
        """
        return {**self.__dict__}

    def __repr__(self):
        def format_value(v):
            if hasattr(v, "shape") and hasattr(v, "dtype"):
                return f"Tensor(dtype={v.dtype}, shape={v.shape})"
            elif isinstance(v, list) and len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
                return f"[Tensor(dtype={v[0].dtype}, shape={v[0].shape}), ...]"
            else:
                return repr(v)

        values_str = "\n".join(f"    {k}: {format_value(v)}" for k, v in self.values.items())
        kwargs_mapping_str = "\n".join(f"    {k}: {v}" for k, v in self.kwargs_mapping.items())

        return f"PipelineState(\n  values={{\n{values_str}\n  }},\n  kwargs_mapping={{\n{kwargs_mapping_str}\n  }}\n)"

mindone.diffusers.modular_pipelines.modular_pipeline.PipelineState.get(keys, default=None)

Get one or multiple values from the pipeline state.

PARAMETER	DESCRIPTION
keys	Key or list of keys for the values TYPE: Union[str, List[str]]
default	The default value to return if not found TYPE: Any DEFAULT: None

RETURNS	DESCRIPTION
Union[Any, Dict[str, Any]]	Union[Any, Dict[str, Any]]: Single value if keys is str, dictionary of values if keys is list

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

def get(self, keys: Union[str, List[str]], default: Any = None) -> Union[Any, Dict[str, Any]]:
    """
    Get one or multiple values from the pipeline state.

    Args:
        keys (Union[str, List[str]]): Key or list of keys for the values
        default (Any): The default value to return if not found

    Returns:
        Union[Any, Dict[str, Any]]: Single value if keys is str, dictionary of values if keys is list
    """
    if isinstance(keys, str):
        return self.values.get(keys, default)
    return {key: self.values.get(key, default) for key in keys}

mindone.diffusers.modular_pipelines.modular_pipeline.PipelineState.get_by_kwargs(kwargs_type)

Get all values with matching kwargs_type.

PARAMETER	DESCRIPTION
kwargs_type	The kwargs_type to filter by TYPE: str

RETURNS	DESCRIPTION
Dict[str, Any]	Dict[str, Any]: Dictionary of values with matching kwargs_type

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

def get_by_kwargs(self, kwargs_type: str) -> Dict[str, Any]:
    """
    Get all values with matching kwargs_type.

    Args:
        kwargs_type (str): The kwargs_type to filter by

    Returns:
        Dict[str, Any]: Dictionary of values with matching kwargs_type
    """
    value_names = self.kwargs_mapping.get(kwargs_type, [])
    return self.get(value_names)

mindone.diffusers.modular_pipelines.modular_pipeline.PipelineState.set(key, value, kwargs_type=None)

Add a value to the pipeline state.

PARAMETER	DESCRIPTION
key	The key for the value TYPE: str
value	The value to store TYPE: Any
kwargs_type	The kwargs_type with which the value is associated TYPE: str DEFAULT: None

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

def set(self, key: str, value: Any, kwargs_type: str = None):
    """
    Add a value to the pipeline state.

    Args:
        key (str): The key for the value
        value (Any): The value to store
        kwargs_type (str): The kwargs_type with which the value is associated
    """
    self.values[key] = value

    if kwargs_type is not None:
        if kwargs_type not in self.kwargs_mapping:
            self.kwargs_mapping[kwargs_type] = [key]
        else:
            self.kwargs_mapping[kwargs_type].append(key)

mindone.diffusers.modular_pipelines.modular_pipeline.PipelineState.to_dict()

Convert PipelineState to a dictionary.

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

def to_dict(self) -> Dict[str, Any]:
    """
    Convert PipelineState to a dictionary.
    """
    return {**self.__dict__}

mindone.diffusers.modular_pipelines.modular_pipeline.BlockState dataclass

Container for block state data with attribute access and formatted representation.

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

@dataclass
class BlockState:
    """
    Container for block state data with attribute access and formatted representation.
    """

    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def __getitem__(self, key: str):
        # allows block_state["foo"]
        return getattr(self, key, None)

    def __setitem__(self, key: str, value: Any):
        # allows block_state["foo"] = "bar"
        setattr(self, key, value)

    def as_dict(self):
        """
        Convert BlockState to a dictionary.

        Returns:
            Dict[str, Any]: Dictionary containing all attributes of the BlockState
        """
        return dict(self.__dict__.items())

    def __repr__(self):
        def format_value(v):
            # Handle tensors directly
            if hasattr(v, "shape") and hasattr(v, "dtype"):
                return f"Tensor(dtype={v.dtype}, shape={v.shape})"

            # Handle lists of tensors
            elif isinstance(v, list):
                if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
                    shapes = [t.shape for t in v]
                    return f"List[{len(v)}] of Tensors with shapes {shapes}"
                return repr(v)

            # Handle tuples of tensors
            elif isinstance(v, tuple):
                if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
                    shapes = [t.shape for t in v]
                    return f"Tuple[{len(v)}] of Tensors with shapes {shapes}"
                return repr(v)

            # Handle dicts with tensor values
            elif isinstance(v, dict):
                formatted_dict = {}
                for k, val in v.items():
                    if hasattr(val, "shape") and hasattr(val, "dtype"):
                        formatted_dict[k] = f"Tensor(shape={val.shape}, dtype={val.dtype})"
                    elif (
                        isinstance(val, list) and len(val) > 0 and hasattr(val[0], "shape") and hasattr(val[0], "dtype")
                    ):
                        shapes = [t.shape for t in val]
                        formatted_dict[k] = f"List[{len(val)}] of Tensors with shapes {shapes}"
                    else:
                        formatted_dict[k] = repr(val)
                return formatted_dict

            # Default case
            return repr(v)

        attributes = "\n".join(f"    {k}: {format_value(v)}" for k, v in self.__dict__.items())
        return f"BlockState(\n{attributes}\n)"

mindone.diffusers.modular_pipelines.modular_pipeline.BlockState.as_dict()

Convert BlockState to a dictionary.

RETURNS	DESCRIPTION
	Dict[str, Any]: Dictionary containing all attributes of the BlockState

Source code in mindone/diffusers/modular_pipelines/modular_pipeline.py

def as_dict(self):
    """
    Convert BlockState to a dictionary.

    Returns:
        Dict[str, Any]: Dictionary containing all attributes of the BlockState
    """
    return dict(self.__dict__.items())