pipelines ¶

@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)
    for stage in self.stages:
        batch = stage(batch, fastvideo_args)

    # 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

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)
        module = PipelineComponentLoader.load_module(
            module_name=load_module_name,
            component_model_path=component_model_path,
            transformers_or_diffusers=transformers_or_diffusers,
            fastvideo_args=fastvideo_args,
        )
        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

def __post_init__(self):
    """Initialize dependent fields after dataclass initialization."""

    # LTX-2 text CFG scales; FLUX uses ``use_embedded_guidance`` so ``guidance_scale > 1`` alone
    # does not enable classifier-free guidance.
    ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
    if self.use_embedded_guidance:
        self.do_classifier_free_guidance = (self.true_cfg_scale > 1.0) or ltx2_text_cfg_enabled
    elif self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
        self.do_classifier_free_guidance = True
    if self.negative_prompt_embeds is None:
        self.negative_prompt_embeds = []
    if self.guidance_scale_2 is None:
        self.guidance_scale_2 = self.guidance_scale

def convert_to_lora_layers(self) -> None:
    """
    Unified method to convert the transformer to a LoRA transformer.
    """
    if self.lora_initialized:
        return
    self.lora_initialized = True
    for (
            transformer_name,
            transformer_module,
    ) in self.trainable_transformer_modules.items():
        converted_count = 0
        # init bookkeeping structures
        if transformer_name not in self.lora_layers:
            # get block list
            block_list = []
            for name, submodule in transformer_module.named_children():
                if isinstance(submodule, nn.ModuleList):
                    block_list = [(f"{name}.{i}", m) for i, m in enumerate(submodule)]
                    break
            self.lora_layers[transformer_name] = LoRAModelLayers(block_list)
        logger.info("Converting %s to LoRA Transformer", transformer_name)
        # scan every module and convert to LoRA layer if applicable

        for block_name, block_modules in _named_module_by_prefix(
                transformer_module,
                list(self.lora_layers[transformer_name].block_mapping),
        ):
            if block_name is not None and (not self.fastvideo_args.training_mode
                                           and self.fastvideo_args.dit_layerwise_offload):
                scope_ctx = _get_hook_ctx(self.lora_layers[transformer_name].block_mapping[block_name])
            else:
                scope_ctx = nullcontext()
            with scope_ctx:
                for name, layer in block_modules:
                    if not self.is_target_layer(name):
                        continue

                    excluded = False
                    for exclude_layer in self.exclude_lora_layers[transformer_name]:
                        if exclude_layer in name:
                            excluded = True
                            break
                    if excluded:
                        continue

                    layer = get_lora_layer(
                        layer,
                        lora_rank=self.lora_rank,
                        lora_alpha=self.lora_alpha,
                        training_mode=self.training_mode,
                    )
                    if layer is not None:
                        block_name_split = name.split(".", 2)
                        if len(block_name_split) > 2:
                            block_name = (block_name_split[0] + "." + block_name_split[1])
                        else:
                            block_name = None
                        if (block_name not in self.lora_layers[transformer_name].block_mapping):
                            block_name = None
                        self.lora_layers[transformer_name].add_lora_layer(block_name, name, layer)
                        replace_submodule(transformer_module, name, layer)
                        converted_count += 1
        logger.info("Converted %d layers to LoRA layers", converted_count)

def set_lora_adapter(self,
                     lora_nickname: str,
                     lora_path: str | None = None,
                     strength: float = 1.0,
                     accumulate: bool = False):  # type: ignore
    """
    Load a LoRA adapter into the pipeline and merge it into the transformer.
    Args:
        lora_nickname: The "nick name" of the adapter when referenced in the pipeline.
        lora_path: The path to the adapter, either a local path or a Hugging Face repo id.
    """

    if lora_nickname not in self.lora_adapters and lora_path is None:
        raise ValueError(f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it.")
    if not self.lora_initialized:
        self.convert_to_lora_layers()
    adapter_updated = False
    rank = dist.get_rank()
    if lora_path is not None and self.lora_adapter_paths.get(lora_nickname) != lora_path:
        self.lora_adapters[lora_nickname] = {}
        lora_local_path = maybe_download_lora(lora_path)
        lora_state_dict = load_file(lora_local_path)

        # Map the hf layer names to our custom layer names
        param_names_mapping_fn = get_param_names_mapping(self.modules["transformer"].param_names_mapping)
        lora_param_names_mapping_fn = get_param_names_mapping(self.modules["transformer"].lora_param_names_mapping)

        # Extract alpha values and weights in a single pass
        to_merge_params: defaultdict[Hashable, dict[Any, Any]] = (defaultdict(dict))
        for name, weight in lora_state_dict.items():
            # Extract weights (lora_A, lora_B, and lora_alpha)
            name = name.replace("diffusion_model.", "")
            name = name.replace(".weight", "")

            if "lora_alpha" in name:
                # Store alpha with minimal mapping - same processing as lora_A/lora_B
                # but store in lora_adapters with ".lora_alpha" suffix
                layer_name = name.replace(".lora_alpha", "")
                layer_name, _, _ = lora_param_names_mapping_fn(layer_name)
                target_name, _, _ = param_names_mapping_fn(layer_name)
                # Store alpha alongside weights with same target_name base
                alpha_key = target_name + ".lora_alpha"
                self.lora_adapters[lora_nickname][alpha_key] = (weight.item()
                                                                if weight.numel() == 1 else float(weight.mean()))
                continue

            name, _, _ = lora_param_names_mapping_fn(name)
            target_name, merge_index, num_params_to_merge = (param_names_mapping_fn(name))
            # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
            # see param mapping in HunyuanVideoArchConfig
            if merge_index is not None and "lora_B" in name:
                to_merge_params[target_name][merge_index] = weight
                if len(to_merge_params[target_name]) == num_params_to_merge:
                    # cat at output dim according to the merge_index order
                    sorted_tensors = [to_merge_params[target_name][i] for i in range(num_params_to_merge)]
                    weight = torch.cat(sorted_tensors, dim=1)
                    del to_merge_params[target_name]
                else:
                    continue

            if target_name in self.lora_adapters[lora_nickname]:
                raise ValueError(f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]")
            self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
        adapter_updated = True
        self.cur_adapter_path = lora_path
        self.lora_adapter_paths[lora_nickname] = lora_path
        logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)

    if (not adapter_updated and self.cur_adapter_name == lora_nickname and self.cur_adapter_strength == strength
            and not accumulate):
        return
    self.cur_adapter_name = lora_nickname
    self.cur_adapter_strength = strength

    # Merge the new adapter
    adapted_count = 0
    for (
            transformer_name,
            transformer_lora_layers,
    ) in self.lora_layers.items():
        for (
                module,
                layers,
        ) in transformer_lora_layers.lora_layers_by_block():
            with _get_hook_ctx(module):
                for name, layer in layers.items():
                    lora_A_name = name + ".lora_A"
                    lora_B_name = name + ".lora_B"
                    lora_alpha_name = name + ".lora_alpha"
                    if (lora_A_name in self.lora_adapters[lora_nickname]
                            and lora_B_name in self.lora_adapters[lora_nickname]):
                        # Get alpha value for this layer (defaults to None if not present)
                        lora_A = self.lora_adapters[lora_nickname][lora_A_name]
                        lora_B = self.lora_adapters[lora_nickname][lora_B_name]
                        # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
                        alpha = self.lora_adapters[lora_nickname].get(lora_alpha_name)
                        try:
                            layer.set_lora_weights(
                                lora_A,
                                lora_B,
                                lora_alpha=alpha,
                                training_mode=self.fastvideo_args.training_mode,
                                lora_path=lora_path,
                                strength=strength,
                                accumulate=accumulate,
                            )
                        except Exception as e:
                            logger.error(
                                "Error setting LoRA weights for layer %s: %s",
                                name,
                                str(e),
                            )
                            raise e
                        adapted_count += 1
                    else:
                        if rank == 0:
                            logger.warning(
                                "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
                                lora_path,
                                name,
                            )
                        layer.disable_lora = True
    logger.info(
        "Rank %d: LoRA adapter %s applied to %d layers",
        rank,
        lora_path,
        adapted_count,
    )