Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

negative_prompt

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</code></pre></details>

Was this page helpful?