Per-rank negative-prompt encoding shared by training model plugins.

Encoding the negative prompt only on rank 0 and broadcasting (the previous Wan path) ran Pipeline.from_pretrained asymmetrically across ranks, which deadlocked on any collective fired during text-encoder load (FSDP device-mesh init, weight broadcast, etc.). The text encoder is small and only loaded once at startup, so loading it on every rank sidesteps the deadlock entirely.

`fastvideo.train.utils.negative_prompt.set_forward_context`

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader`

TextEncoderLoader(device = None)

Loader for text encoders.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.Source`

Source(model_or_path: str, prefix: str = '', fall_back_to_pt: bool = True, allow_patterns_overrides: list[str] | None = None)

A source for weights.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.Source.model_or_path`

model_or_path: str

The model ID or path.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.Source.prefix`

prefix: str = ''

A prefix to prepend to all weights.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.Source.fall_back_to_pt`

fall_back_to_pt: bool = True

Whether .pt weights can be used.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.Source.allow_patterns_overrides`

allow_patterns_overrides: list[str] | None = None

If defined, weights will load exclusively using these patterns.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.load`

load(model_path: str, fastvideo_args: FastVideoArgs)

Load the text encoders based on the model path, and inference args.

`fastvideo.train.utils.negative_prompt.TextEncoderLoader.load_model`

load_model(model_path: str, model_config: EncoderConfig, target_device: torch.device, fastvideo_args: FastVideoArgs, dtype: str = 'fp16', use_text_encoder_override: bool = False, cpu_offload: bool | None = None, offload_flag: str = 'text_encoder_cpu_offload')

`fastvideo.train.utils.negative_prompt.make_inference_args`

make_inference_args(tc: TrainingConfig, *, model_path: str) -> TrainingArgs

Build a TrainingArgs for inference (validation / pipelines).

`fastvideo.train.utils.negative_prompt.maybe_download_model`

maybe_download_model(model_name_or_path: str, local_dir: str | None = None, download: bool = True, revision: str | None = None, allow_patterns: list[str] | None = None) -> str

Check if the model path is a Hugging Face Hub model ID and download it if needed.

Supports an "umbrella" repo layout where a single HF repo holds multiple pipeline variants under sibling subfolders. If the input is shaped as org/repo/subfolder (i.e. a non-existent local path with 3+ slash- separated components and at least one segment that does not look like a posix-absolute path), treat the first two components as the HF repo id and the remainder as a subfolder; only the subfolder's blobs are downloaded, and the returned local path points inside that subfolder.

Parameters

model_name_or_path: str: Local path, Hugging Face Hub model ID, or org/repo/subfolder umbrella-repo reference.
local_dir: str | None: Local directory to save the model
download: bool: Whether to download the model from Hugging Face Hub
revision: str | None: Optional immutable Hub revision.
allow_patterns: list[str] | None: Optional Hub glob patterns limiting downloaded files. Local paths are returned unchanged. For umbrella references, the patterns are interpreted relative to the selected subfolder.

Returns

str: Local path to the model (or to the subfolder inside the snapshot).

`fastvideo.train.utils.negative_prompt.TrainingConfig`

TrainingConfig(distributed: DistributedConfig = DistributedConfig(), data: DataConfig = DataConfig(), optimizer: OptimizerConfig = OptimizerConfig(), loop: TrainingLoopConfig = TrainingLoopConfig(), checkpoint: CheckpointConfig = CheckpointConfig(), tracker: TrackerConfig = TrackerConfig(), performance: PerformanceConfig = PerformanceConfig(), vsa_sparsity: float = 0.0, vsa_cache_tile_buf: bool = False, model: ModelTrainingConfig = ModelTrainingConfig(), pipeline_config: PipelineConfig | None = None, model_path: str = '', dit_precision: str = 'fp32', vsa_ref_keep_rate: float | None = None, vsa_tile_size: int = 256)

`fastvideo.train.utils.negative_prompt.encode_negative_prompt`

encode_negative_prompt(training_config: TrainingConfig, *, prompt: str, device: torch.device, dtype: torch.dtype, encoder_index: int = 0) -> tuple[torch.Tensor, torch.Tensor]

Per-rank encode of prompt using encoder encoder_index.

Reads pipeline_config.text_encoder_configs[encoder_index] so the encoder class (e.g. UMT5 for Wan) and tokenizer kwargs match the inference path, and applies the matching postprocess_text_funcs entry. Returns (embeds, mask) on device cast to dtype.

Source

def encode_negative_prompt(
    training_config: TrainingConfig,
    *,
    prompt: str,
    device: torch.device,
    dtype: torch.dtype,
    encoder_index: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Per-rank encode of ``prompt`` using encoder ``encoder_index``.

    Reads ``pipeline_config.text_encoder_configs[encoder_index]`` so the
    encoder class (e.g. UMT5 for Wan) and tokenizer kwargs match the
    inference path, and applies the matching ``postprocess_text_funcs``
    entry. Returns ``(embeds, mask)`` on ``device`` cast to ``dtype``.
    """
    tc = training_config
    pipeline_config = tc.pipeline_config
    if pipeline_config is None:
        raise ValueError("training_config.pipeline_config is required for negative "
                         "prompt encoding")

    encoder_configs = pipeline_config.text_encoder_configs
    postprocess_funcs = pipeline_config.postprocess_text_funcs
    preprocess_funcs = getattr(pipeline_config, "preprocess_text_funcs", None)

    if encoder_index < 0 or encoder_index >= len(encoder_configs):
        raise IndexError(f"encoder_index {encoder_index} out of range for "
                         f"text_encoder_configs (len={len(encoder_configs)})")
    encoder_config = encoder_configs[encoder_index]
    postprocess_text = postprocess_funcs[encoder_index]
    preprocess_text = (preprocess_funcs[encoder_index] if preprocess_funcs is not None else None)

    # HF convention: text_encoder / tokenizer for index 0,
    # text_encoder_2 / tokenizer_2 for index 1, etc.
    suffix = "" if encoder_index == 0 else f"_{encoder_index + 1}"
    encoder_subdir = f"text_encoder{suffix}"
    tokenizer_subdir = f"tokenizer{suffix}"

    model_path = maybe_download_model(tc.model_path)
    inference_args = make_inference_args(tc, model_path=model_path)
    # Keep the encoder on-device; CPU offload would init an FSDP device
    # mesh and reintroduce the collective at load time.
    inference_args.text_encoder_cpu_offload = False

    loader = TextEncoderLoader()
    text_encoder = loader.load(
        os.path.join(model_path, encoder_subdir),
        inference_args,
    ).to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, tokenizer_subdir))

    tok_kwargs = dict(encoder_config.tokenizer_kwargs)
    text = preprocess_text(prompt) if preprocess_text is not None else prompt

    with torch.no_grad(), set_forward_context(
            current_timestep=0,
            attn_metadata=None,
    ):
        text_inputs = tokenizer(text, **tok_kwargs).to(device)
        outputs = text_encoder(
            input_ids=text_inputs.input_ids,
            attention_mask=text_inputs.attention_mask,
        )
        # Mirror TextEncodingStage: postprocess reads outputs.attention_mask.
        outputs.attention_mask = text_inputs["attention_mask"]
        embeds = postprocess_text(outputs).to(device=device, dtype=dtype)
        mask = text_inputs["attention_mask"].to(device=device, dtype=dtype)

    del text_encoder, tokenizer

    return embeds, mask