-
Notifications
You must be signed in to change notification settings - Fork 2
Advanced
Pipeline components can be configured in workflow JSON files to customize their behavior and performance characteristics. Components can be configured at both the pipeline level and individual component level.
The following components can be configured, assuming that they are used by a given pipeline:
- transformer
- vae (Variational Autoencoder)
- unet
- text_encoder
- text_encoder_2
- text_encoder_3
- tokenizer
- tokenizer_2
- tokenizer_3
- image_encoder
- feature_extractor
- model
- controlnet
Each component follows this basic structure:
{
"configuration": {
"component_type": "required.module.type"
},
"from_pretrained_arguments": {
"model_name": "model_id_or_path",
// other arguments...
},
"quantization_config": {
"configuration": {
"config_type": "quantization_type"
},
"arguments": {
// quantization specific arguments
}
}
}-
configuration.component_type: The Python class path for the component -
from_pretrained_arguments: Arguments for loading the pretrained model
Here's an example of configuring a VAE component with quantization:
"vae": {
"configuration": {
"component_type": "AutoencoderKLCogVideoX"
},
"quantization_config": {
"configuration": {
"config_type": "TorchAoConfig"
},
"arguments": {
"quant_type": "{int4wo}"
}
},
"from_pretrained_arguments": {
"model_name": "model_name_here",
"subfolder": "vae",
"torch_dtype": "torch.bfloat16"
}
}The main pipeline can be configured with these options:
"configuration": {
"component_type": "required.pipeline.type",
"offload": "model|sequential",
"no_generator": true|false,
"do_not_send_to_device": true|false,
"vae": {
"enable_slicing": true|false,
"enable_tiling": true|false,
"set_memory_format": true|false
},
"unet": {
"enable_forward_chunking": true|false,
"set_memory_format": true|false
},
"shared_components": ["component_names"],
"reused_components": ["component_names"]
}-
offload: Controls CPU offloading strategy-
model: Enables model CPU offloading -
sequential: Enables sequential CPU offloading
-
-
no_generator: Disables random generator for pipelines that don't support it - vae settings:
-
enable_slicing: Enables VAE slicing for memory efficiency -
enable_tiling: Enables VAE tiling for large images -
set_memory_format: Sets channels_last memory format
-
- unet settings:
-
enable_forward_chunking: Enables forward pass chunking -
set_memory_format: Sets channels_last memory format
-
- Component sharing:
-
shared_components: List of components to share with other pipelines -
reused_components: List of components to reuse from previous pipelines
-
- Use CPU offloading when working with large models that don't fit in GPU memory
- Enable VAE slicing/tiling for processing large images
- Share components between pipeline steps when possible to reduce memory usage
- Consider quantization for large models to reduce memory footprint
- Use appropriate memory formats for your use case
The quantization configuration allows you to specify how model components should be quantized to reduce memory usage and potentially improve performance. This is defined in the workflow JSON schema and processed by the pipeline system.
The quantization config can be specified for pipeline components like transformers, text encoders, etc. Here's the relevant schema structure:
"quantization_config": {
"type": "object",
"properties": {
"configuration": {
"type": "object",
"properties": {
"config_type": {
"type": "string"
}
},
"required": ["config_type"]
},
"arguments": {
"$ref": "#/$defs/arguments"
}
},
"required": ["configuration", "arguments"]
}Here's an example of using quantization config in a workflow file:
"transformer": {
"configuration": {
"component_type": "SD3Transformer2DModel"
},
"quantization_config": {
"configuration": {
"config_type": "BitsAndBytesConfig"
},
"arguments": {
"load_in_4bit": true,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": "torch.bfloat16"
}
},
"from_pretrained_arguments": {
"model_name": "stabilityai/stable-diffusion-3.5-large",
"subfolder": "transformer",
"torch_dtype": "torch.bfloat16"
}
}The BitsAndBytesConfig is used for 4-bit and 8-bit quantization using the bitsandbytes library. Common arguments include:
-
load_in_4bit: Enable 4-bit quantization -
bnb_4bit_quant_type: Quantization type (e.g., "nf4") -
bnb_4bit_compute_dtype: Computation data type
The TorchAoConfig is used for 4-bit and 8-bit quantization using the torchao library. Common arguments include:
-
quant_type: The type of quantization we want to use (e.g., "int4wo", "int4dq", "int8wo", "int8dq")
The GGUFQuantizationConfig is used for 4-bit and 8-bit quantization using the gguf file format. Common arguments include:
-
compute_dtype: This sets the computational type which might be different than the input type (e.g., torch.float32)
- Quantization config is optional for pipeline components
- The
config_typeshould be a valid configuration class (typically from transformers or diffusers) - Arguments are passed directly to the
config_typeconstructor - Improper configuration may result in runtime errors when loading the model