Spaces:
Running on Zero
Running on Zero
| """ | |
| Communication module for distributed inference pipeline. | |
| This module provides abstractions for distributed communication operations, | |
| model data transfer, and buffer management in the StreamDiffusionV2 pipeline. | |
| """ | |
| from .distributed_communicator import DistributedCommunicator | |
| from .model_data_transfer import ModelDataTransfer | |
| from .buffer_manager import BufferManager | |
| from .data_containers import LatentData, KVCacheData, CommunicationConfig | |
| from .kv_cache_manager import KVCacheManager | |
| from .utils import CommunicationTags, init_distributed, setup_logging, compute_balanced_split | |
| __all__ = [ | |
| 'DistributedCommunicator', | |
| 'ModelDataTransfer', | |
| 'BufferManager', | |
| 'LatentData', | |
| 'KVCacheData', | |
| 'CommunicationConfig', | |
| 'KVCacheManager', | |
| 'CommunicationTags', | |
| 'init_distributed', | |
| 'setup_logging', | |
| 'compute_balanced_split' | |
| ] | |