File size: 2,095 Bytes
a8fae71 6c349c3 a8fae71 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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,
# repo="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 |