YiYiXu HF Staff

Update block.py

2aca20b verified 9 months ago

2.12 kB

from typing import List

from diffusers.modular_pipelines import (
    ComponentSpec,
    InputParam,
    ModularPipelineBlocks,
    OutputParam,
    PipelineState,
)


class MyCustomBlock(ModularPipelineBlocks):
    """
    A custom block for [describe what your block does].
    
    Replace this with your implementation.
    """

    @property
    def description(self) -> str:
        """Description of the block."""
        return "A template custom block - replace with your description"

    
    @property
    def expected_components(self) -> List[ComponentSpec]:
        """Define model components your block needs (e.g., transformers, VAEs)."""
        return [
            # Example:
            # ComponentSpec(
            #     name="model",
            #     type_hint=SomeModelClass,
            #     pretrained_model_name_or_path="organization/model-name",
            # ),
        ]

    @property
    def inputs(self) -> List[InputParam]:
        """Define input parameters for your block."""
        return [
            # Example:
            # InputParam(
            #     "prompt",
            #     type_hint=str,
            #     required=True,
            #     description="Input prompt",
            #     metadata={"mellon": "textbox"},  # For Mellon UI
            # ),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        """Define output parameters for your block."""
        return [
            # Example:
            # OutputParam(
            #     "result",
            #     type_hint=str,
            #     description="Output result",
            #     metadata={"mellon": "text"},  # For Mellon UI
            # ),
        ]

    def __call__(self, components, state: PipelineState) -> PipelineState:
        """Execute your block logic."""
        block_state = self.get_block_state(state)
        
        # Your implementation here
        # Access inputs via block_state.<input_name>
        # Set outputs via block_state.<output_name> = value
        
        self.set_block_state(state, block_state)
        return components, state

	from typing import List

	from diffusers.modular_pipelines import (
	ComponentSpec,
	InputParam,
	ModularPipelineBlocks,
	OutputParam,
	PipelineState,
	)


	class MyCustomBlock(ModularPipelineBlocks):
	"""
	A custom block for [describe what your block does].

	Replace this with your implementation.
	"""

	@property
	def description(self) -> str:
	"""Description of the block."""
	return "A template custom block - replace with your description"


	@property
	def expected_components(self) -> List[ComponentSpec]:
	"""Define model components your block needs (e.g., transformers, VAEs)."""
	return [
	# Example:
	# ComponentSpec(
	# name="model",
	# type_hint=SomeModelClass,
	# pretrained_model_name_or_path="organization/model-name",
	# ),
	]

	@property
	def inputs(self) -> List[InputParam]:
	"""Define input parameters for your block."""
	return [
	# Example:
	# InputParam(
	# "prompt",
	# type_hint=str,
	# required=True,
	# description="Input prompt",
	# metadata={"mellon": "textbox"}, # For Mellon UI
	# ),
	]

	@property
	def intermediate_outputs(self) -> List[OutputParam]:
	"""Define output parameters for your block."""
	return [
	# Example:
	# OutputParam(
	# "result",
	# type_hint=str,
	# description="Output result",
	# metadata={"mellon": "text"}, # For Mellon UI
	# ),
	]

	def __call__(self, components, state: PipelineState) -> PipelineState:
	"""Execute your block logic."""
	block_state = self.get_block_state(state)

	# Your implementation here
	# Access inputs via block_state.<input_name>
	# Set outputs via block_state.<output_name> = value

	self.set_block_state(state, block_state)
	return components, state