streamv2v/communication/__init__.py

"""
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'
]