from abc import ABC, abstractmethod import torch from fastvideo.v1.inference_args import InferenceArgs from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler from fastvideo.v1.logger import init_logger from fastvideo.v1.platforms import current_platform from fastvideo.v1.distributed.parallel_state import get_sp_group from fastvideo.v1.pipelines.composed import ComposedPipelineBase import os import json logger = init_logger(__name__) T2V_PIPLINES = { "HunyuanVideo": ("hunyuan_video", "HunyuanVideoPipeline"), } FASTVIDEO_PIPLINES = { **T2V_PIPLINES, } def resolve_pipeline_cls(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name from fastvideo.v1.pipelines.hunyuan.hunyuan_video import HunyuanVideoPipeline return HunyuanVideoPipeline def resolve_pipeline_cls_v2(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name from fastvideo.v1.pipelines.implementations.hunyuan import HunyuanVideoPipeline return HunyuanVideoPipeline class ModelComponents: """Container for model components used by diffusion pipelines""" def __init__(self, vae=None, text_encoder=None, text_encoder_2=None, transformer=None, scheduler=None): self.vae = vae self.text_encoder = text_encoder self.text_encoder_2 = text_encoder_2 self.transformer = transformer self.scheduler = scheduler HF_PIPELINE_CONFIG = \ """ { "_class_name": "HunyuanVideoPipeline", "_diffusers_version": "0.32.0.dev0", "scheduler": [ "diffusers", "FlowMatchEulerDiscreteScheduler" ], "text_encoder": [ "transformers", "LlamaModel" ], "text_encoder_2": [ "transformers", "CLIPTextModel" ], "tokenizer": [ "transformers", "LlamaTokenizerFast" ], "tokenizer_2": [ "transformers", "CLIPTokenizer" ], "transformer": [ "diffusers", "HunyuanVideoTransformer3DModel" ], "vae": [ "diffusers", "AutoencoderKLHunyuanVideo" ] } """ class PipelineLoader(ABC): """ Base class for loading model components that can be used across different pipelines. This separates model loading from pipeline creation and inference. """ def __init__(self, inference_args: InferenceArgs, use_v0=False): assert current_platform.device_type == "cuda" if use_v0: import os local_rank = int(os.environ.get("LOCAL_RANK", 0)) logger.info(f"Using device: cuda:{local_rank}") self.device = torch.device(f"cuda:{local_rank}") else: sp_group = get_sp_group() # device = "cuda" if sp_group.world_size > 1: local_rank = sp_group.rank logger.info(f"Using device: cuda:{local_rank}") self.device = torch.device(f"cuda:{local_rank}") @abstractmethod def load_text_encoders(self, inference_args: InferenceArgs): """Load the text encoders based on the inference args.""" raise NotImplementedError @abstractmethod def load_vae(self, inference_args: InferenceArgs): """Load the VAE based on the inference args.""" raise NotImplementedError @abstractmethod def load_transformer(self, inference_args: InferenceArgs): """Load the transformer based on the inference args.""" raise NotImplementedError @abstractmethod def load_scheduler(self, inference_args: InferenceArgs): """Load the scheduler based on the inference args.""" raise NotImplementedError def load_components(self, inference_args: InferenceArgs): logger.info(f"Loading components for {inference_args.model_path}") logger.info("Loading transformer") transformer = self.load_transformer(inference_args) logger.info("Loading VAE") vae, vae_kwargs = self.load_vae(inference_args) logger.info("Loading text encoders") text_encoders = self.load_text_encoders(inference_args) logger.info(f"\tLoaded {len(text_encoders)} text encoders") logger.info(f"\ttext_encoders: {[type(e) for e in text_encoders]}") text_encoder = text_encoders[0] text_encoder_2 = text_encoders[1] # may be None logger.info("Loading scheduler") scheduler = self.load_scheduler(inference_args) logger.info("Model components loaded") return vae, vae_kwargs, text_encoder, text_encoder_2, transformer, scheduler def load_pipeline_v2(self, inference_args: InferenceArgs): """Load the staged pipeline based on the inference args.""" # Load configuration # TODO(will): load config from model path. also move to better place. pipeline_cls = resolve_pipeline_cls_v2(inference_args) assert issubclass(pipeline_cls, ComposedPipelineBase) logger.info(f"Pipeline class: {pipeline_cls}") config = json.loads(HF_PIPELINE_CONFIG) logger.info(f"Config: {config}") logger.info(f"Loading components for {inference_args.model_path}") logger.info("Loading transformer") transformer = self.load_transformer(inference_args) logger.info("Loading VAE") vae, vae_kwargs = self.load_vae(inference_args) logger.info("Loading text encoders") text_encoders = self.load_text_encoders(inference_args) logger.info(f"\tLoaded {len(text_encoders)} text encoders") logger.info(f"\ttext_encoders: {[type(e) for e in text_encoders]}") text_encoder = text_encoders[0] text_encoder_2 = text_encoders[1] # may be None logger.info("Loading scheduler") scheduler = self.load_scheduler(inference_args) # Create pipeline pipeline = pipeline_cls() # TODO(will): the keys of this dict should be from the hf config pipeline_modules = { "vae": vae, "text_encoder": text_encoder, "text_encoder_2": text_encoder_2, "transformer": transformer, "scheduler": scheduler } logger.info(f"Registering modules") pipeline.register_modules(pipeline_modules) logger.info(f"Setting up pipeline") pipeline.setup_pipeline(inference_args) logger.info(f"Initializing pipeline") pipeline.initialize_pipeline(inference_args) return pipeline def load_pipeline(self, inference_args: InferenceArgs): """Load the pipeline based on the inference args.""" # Load configuration # TODO(will): load config from model path. also move to better place. pipeline_cls = resolve_pipeline_cls(inference_args) logger.info(f"Pipeline class: {pipeline_cls}") logger.info(f"Loading components for {inference_args.model_path}") logger.info("Loading transformer") transformer = self.load_transformer(inference_args) logger.info("Loading VAE") vae, vae_kwargs = self.load_vae(inference_args) logger.info("Loading text encoders") text_encoders = self.load_text_encoders(inference_args) logger.info(f"\tLoaded {len(text_encoders)} text encoders") logger.info(f"\ttext_encoders: {[type(e) for e in text_encoders]}") text_encoder = text_encoders[0] text_encoder_2 = text_encoders[1] # may be None logger.info("Loading scheduler") scheduler = self.load_scheduler(inference_args) logger.info("Model components loaded") # Create pipeline pipeline = pipeline_cls( vae=vae, text_encoder=text_encoder, text_encoder_2=text_encoder_2, transformer=transformer, scheduler=scheduler ) return pipeline def load_scheduler(self, inference_args: InferenceArgs): """Create a scheduler based on the inference args. Can be overridden by subclasses.""" if hasattr(inference_args, 'denoise_type') and inference_args.denoise_type == "flow": return FlowMatchDiscreteScheduler( shift=inference_args.flow_shift, reverse=inference_args.flow_reverse, solver=inference_args.flow_solver, ) return None @staticmethod def load_pipeline_from_config(config: dict): """Load a pipeline from a config.""" pipeline_cls = resolve_pipeline_cls(inference_args) assert issubclass(pipeline_cls, ComposedPipelineBase) logger.info(f"Pipeline class: {pipeline_cls}") def get_pipeline_loader(inference_args: InferenceArgs) -> PipelineLoader: """Get a pipeline loader based on the inference args.""" from fastvideo.v1.pipelines.implementations.hunyuan.hunyuan_loader import HunyuanPipelineLoader # Map model types to their specific loaders loader_map = { "hunyuan_video": HunyuanPipelineLoader, # Add other model types and their loaders here } # Determine which loader to use based on model type # model_type = inference_args.model_type if hasattr(inference_args, 'model_type') else None # if model_type and model_type in loader_map: # return loader_map[model_type]() # TODO(will): add more model types and their loaders here instead of # hardcoding return HunyuanPipelineLoader(inference_args, use_v0=True) # Default to base loader return PipelineLoader()