composed_pipeline_base ¶

@abstractmethod
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
    """
    Create the inference pipeline stages.
    """
    raise NotImplementedError

def create_training_stages(self, training_args: TrainingArgs):
    """
    Create the training pipeline stages.
    """
    raise NotImplementedError

@torch.no_grad()
def forward(
    self,
    batch: ForwardBatch,
    fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
    """
    Generate a video or image using the pipeline.

    Args:
        batch: The batch to generate from.
        fastvideo_args: The inference arguments.
    Returns:
        ForwardBatch: The batch with the generated video or image.
    """
    if not self.post_init_called:
        self.post_init()

    # Execute each stage
    logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys())
    # logger.info("Batch: %s", batch)
    try:
        for stage in self.stages:
            batch = stage(batch, fastvideo_args)
    except BaseException:
        # A stage's own hook frees only what that stage was the last user
        # of. When the run aborts earlier, everything already materialized
        # stays for the life of the generator, and the retry a
        # memory-constrained caller is most likely to attempt starts from a
        # worse position than the request that just failed.
        self._release_all_lazy_modules()
        raise

    # Return the output
    return batch

@classmethod
def from_pretrained(cls,
                    model_path: str,
                    device: str | None = None,
                    torch_dtype: torch.dtype | None = None,
                    pipeline_config: str | PipelineConfig | None = None,
                    args: argparse.Namespace | None = None,
                    required_config_modules: list[str] | None = None,
                    loaded_modules: dict[str, torch.nn.Module]
                    | None = None,
                    **kwargs) -> "ComposedPipelineBase":
    """
    Load a pipeline from a pretrained model.
    loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
    If provided, loaded_modules will be used instead of loading from config/pretrained weights.
    """
    if args is None or args.inference_mode:

        kwargs['model_path'] = model_path
        fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
    else:
        assert args is not None, "args must be provided for training mode"
        fastvideo_args = TrainingArgs.from_cli_args(args)
        # TODO(will): fix this so that its not so ugly
        fastvideo_args.model_path = model_path
        for key, value in kwargs.items():
            setattr(fastvideo_args, key, value)

        fastvideo_args.dit_cpu_offload = False
        # we hijack the precision to be the master weight type so that the
        # model is loaded with the correct precision. Subsequently we will
        # use FSDP2's MixedPrecisionPolicy to set the precision for the
        # fwd, bwd, and other operations' precision.
        assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'

    logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)

    pipe = cls(model_path,
               fastvideo_args,
               required_config_modules=required_config_modules,
               loaded_modules=loaded_modules)
    pipe.post_init()
    return pipe

@classmethod
def get_hf_download_allow_patterns(cls) -> list[str] | None:
    """Return Hub patterns for the manifest and selected components."""
    component_dirs = cls.get_hf_download_component_dirs()
    if component_dirs is None:
        return None
    return [
        "model_index.json",
        "modular_model_index.json",
        *(f"{component_dir}/**" for component_dir in component_dirs),
    ]

@classmethod
def get_hf_download_component_dirs(cls) -> tuple[str, ...] | None:
    """Return component directories for an opt-in partial Hub download."""
    return None

def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
    """
    Initialize the pipeline.
    """
    return

def load_modules(self,
                 fastvideo_args: FastVideoArgs,
                 loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]:
    """
    Load the modules from the config.
    loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, 
    If provided, loaded_modules will be used instead of loading from config/pretrained weights.
    """

    model_index = self._load_config(self.model_path)
    logger.info("Loading pipeline modules from config: %s", model_index)

    # remove keys that are not pipeline modules
    model_index.pop("_class_name")
    model_index.pop("_diffusers_version")
    model_index.pop("_name_or_path", None)
    model_index.pop("workload_type", None)
    if "boundary_ratio" in model_index and model_index["boundary_ratio"] is not None:
        logger.info("MoE pipeline detected. Adding transformer_2 to self.required_config_modules...")
        self.required_config_modules.append("transformer_2")
        logger.info("MoE pipeline detected. Setting boundary ratio to %s", model_index["boundary_ratio"])
        fastvideo_args.pipeline_config.dit_config.boundary_ratio = model_index["boundary_ratio"]

    model_index.pop("boundary_ratio", None)
    # used by Wan2.2 ti2v
    model_index.pop("expand_timesteps", None)
    # HF metadata (e.g. Flux2 Klein is_distilled); not a loadable module
    model_index.pop("is_distilled", None)

    # some sanity checks
    assert len(model_index) > 1, "model_index.json must contain at least one pipeline module"

    for module_name in self.required_config_modules:
        if module_name not in model_index and module_name in self._extra_config_module_map:
            extra_module_value = self._extra_config_module_map[module_name]
            logger.warning(
                "model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.",
                module_name, module_name, extra_module_value)
            if extra_module_value in model_index:
                logger.info("Using module %s for %s", extra_module_value, module_name)
                model_index[module_name] = model_index[extra_module_value]
                continue
            else:
                raise ValueError(
                    f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}"
                )

    # all the component models used by the pipeline
    required_modules = self.required_config_modules
    logger.info("Loading required modules: %s", required_modules)

    modules = {}
    for module_name, module_spec in model_index.items():
        if not isinstance(module_spec, list | tuple):
            logger.info(
                "Skipping non-module config entry %s=%s",
                module_name,
                module_spec,
            )
            continue
        if len(module_spec) < 1:
            logger.warning(
                "Skipping module %s due to invalid empty spec in model_index.json",
                module_name,
            )
            continue
        transformers_or_diffusers = module_spec[0]
        if transformers_or_diffusers is None:
            logger.warning("Module %s in model_index.json has null value, removing from required_config_modules",
                           module_name)
            if module_name in self.required_config_modules:
                self.required_config_modules.remove(module_name)
            continue
        if module_name not in required_modules:
            logger.info("Skipping module %s", module_name)
            continue
        if loaded_modules is not None and module_name in loaded_modules:
            logger.info("Using module %s already provided", module_name)
            modules[module_name] = loaded_modules[module_name]
            continue

        # we load the module from the extra config module map if it exists
        if module_name in self._extra_config_module_map:
            load_module_name = self._extra_config_module_map[module_name]
        else:
            load_module_name = module_name

        component_model_path = os.path.join(self.model_path, load_module_name)

        def load_component(load_module_name: str = load_module_name,
                           component_model_path: str = component_model_path,
                           transformers_or_diffusers: str = transformers_or_diffusers) -> Any:
            return PipelineComponentLoader.load_module(
                module_name=load_module_name,
                component_model_path=component_model_path,
                transformers_or_diffusers=transformers_or_diffusers,
                fastvideo_args=fastvideo_args,
            )

        if self._lazy_module_load_enabled(fastvideo_args) and module_name in self._lazy_module_names:
            module = LazyModule(module_name, load_component)
            logger.info("Deferred module %s from %s", module_name, component_model_path)
        else:
            module = load_component()
            logger.info("Loaded module %s from %s", module_name, component_model_path)

        if module_name in modules:
            logger.warning("Overwriting module %s", module_name)
        modules[module_name] = module

    # Check if all required modules were loaded
    for module_name in required_modules:
        if module_name not in modules or modules[module_name] is None:
            raise ValueError(
                f"Required module key: {module_name} value: {modules.get(module_name)} was not found in loaded modules {modules.keys()}"
            )

    return modules