--- title: reason1 ---

Reason1 (Qwen2.5-VL) text encoder.

fastvideo.models.encoders.reason1.BaseEncoderOutput

BaseEncoderOutput(last_hidden_state: torch.FloatTensor | None = None, pooler_output: torch.FloatTensor | None = None, hidden_states: tuple[torch.FloatTensor, ...] | None = None, attentions: tuple[torch.FloatTensor, ...] | None = None, attention_mask: torch.Tensor | None = None)

fastvideo.models.encoders.reason1.Reason1Config

Reason1Config(arch_config: Reason1ArchConfig = Reason1ArchConfig(), prefix: str = '', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, is_chat_model: bool = False, treat_empty_as_dot: bool = False, tokenizer_type: str = 'Qwen/Qwen2.5-VL-7B-Instruct', *, _resolved_attention_backend: AttentionBackendEnum | None = None, chat_template_enable_thinking: bool = False)

Reason1 text encoder config.

fastvideo.models.encoders.reason1.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.encoders.reason1.TextEncoder

TextEncoder(config: TextEncoderConfig)

Base for native encoders with a model-specific forward output contract.

fastvideo.models.encoders.reason1.TextEncoder.forward

forward(*args: Any, **kwargs: Any) -> TextEncoderOutputT

fastvideo.models.encoders.reason1.default_weight_loader

default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None

Default weight loader.

fastvideo.models.encoders.reason1.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple

Qwen2_5_VLForConditionalGenerationSimple(config)

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.init_weights

init_weights(buffer_device: Optional[torch.device] = None)

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.can_generate

can_generate()

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.get_input_embeddings

get_input_embeddings()

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.set_input_embeddings

set_input_embeddings(value)

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.get_output_embeddings

get_output_embeddings()

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.set_output_embeddings

set_output_embeddings(new_embeddings)

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.set_decoder

set_decoder(decoder)

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.get_decoder

get_decoder()

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.get_rope_index

get_rope_index(input_ids: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]

Calculate the 3D rope index based on image and video's temporal, height and width in LLM.

Explanation

Each embedding sequence contains vision embedding and text embedding or just contains text embedding.

For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. Examples: input_ids: [T T T T T], here T is for text. temporal position_ids: [0, 1, 2, 3, 4] height position_ids: [0, 1, 2, 3, 4] width position_ids: [0, 1, 2, 3, 4]

For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part and 1D rotary position embeddin for text part. Examples: Temporal (Time): 3 patches, representing different segments of the video in time. Height: 2 patches, dividing each frame vertically. Width: 2 patches, dividing each frame horizontally. We also have some important parameters: fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] text temporal position_ids: [101, 102, 103, 104, 105] text height position_ids: [101, 102, 103, 104, 105] text width position_ids: [101, 102, 103, 104, 105] Here we calculate the text start position_ids as the max vision position_ids plus 1.

Parameters

input_ids: `torch.LongTensor` of shape `(batch_size, sequence_length)`

Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it.

image_grid_thw: `torch.LongTensor` of shape `(num_images, 3)`, *optional*

The temporal, height and width of feature shape of each image in LLM.

video_grid_thw: `torch.LongTensor` of shape `(num_videos, 3)`, *optional*

The temporal, height and width of feature shape of each video in LLM.

second_per_grid_ts: `torch.Tensor` of shape `(num_videos)`, *optional*

The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.

attention_mask: `torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*

Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:

  • 1 for tokens that are not masked,
  • 0 for tokens that are masked.

Returns

torch.Tensor

position_ids (torch.LongTensor of shape (3, batch_size, sequence_length))

torch.Tensor

mrope_position_deltas (torch.Tensor of shape (batch_size))

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.forward

forward(input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, pixel_values: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]

Parameters

labels: `torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*

Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].

Returns:

Example:

>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

>>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
>>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

>>> messages = [
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": "What is shown in this image?"},
        ],
    },
]
>>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
>>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."

fastvideo.models.encoders.reason1.Qwen2_5_VLForConditionalGenerationSimple.prepare_inputs_for_generation

prepare_inputs_for_generation(input_ids, past_key_values = None, attention_mask = None, inputs_embeds = None, cache_position = None, position_ids = None, use_cache = True, pixel_values = None, pixel_values_videos = None, image_grid_thw = None, video_grid_thw = None, second_per_grid_ts = None, **kwargs)

fastvideo.models.encoders.reason1.get_rope_index

get_rope_index(model_config, input_ids: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]

Calculate the 3D rope index based on image and video's temporal, height and width in LLM.

Explanation

Each embedding sequence contains vision embedding and text embedding or just contains text embedding.

For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. Examples: input_ids: [T T T T T], here T is for text. temporal position_ids: [0, 1, 2, 3, 4] height position_ids: [0, 1, 2, 3, 4] width position_ids: [0, 1, 2, 3, 4]

For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part and 1D rotary position embeddin for text part. Examples: Temporal (Time): 3 patches, representing different segments of the video in time. Height: 2 patches, dividing each frame vertically. Width: 2 patches, dividing each frame horizontally. We also have some important parameters: fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] text temporal position_ids: [101, 102, 103, 104, 105] text height position_ids: [101, 102, 103, 104, 105] text width position_ids: [101, 102, 103, 104, 105] Here we calculate the text start position_ids as the max vision position_ids plus 1.

Parameters

input_ids: `torch.LongTensor` of shape `(batch_size, sequence_length)`

Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it.

image_grid_thw: `torch.LongTensor` of shape `(num_images, 3)`, *optional*

The temporal, height and width of feature shape of each image in LLM.

video_grid_thw: `torch.LongTensor` of shape `(num_videos, 3)`, *optional*

The temporal, height and width of feature shape of each video in LLM.

second_per_grid_ts: `torch.Tensor` of shape `(num_videos)`, *optional*

The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.

attention_mask: `torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*

Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:

  • 1 for tokens that are not masked,
  • 0 for tokens that are masked.

Returns

torch.Tensor

position_ids (torch.LongTensor of shape (3, batch_size, sequence_length))

torch.Tensor

mrope_position_deltas (torch.Tensor of shape (batch_size))

fastvideo.models.encoders.reason1.Reason1TextEncoder

Reason1TextEncoder(config: Reason1Config, prefix: str = '', checkpoint_path: str | None = None)

Reason1 (Qwen2.5-VL) text encoder.

Source
class Reason1TextEncoder(TextEncoder):
    """Reason1 (Qwen2.5-VL) text encoder."""

    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.TORCH_SDPA,
    )

    def __init__(self, config: Reason1Config, prefix: str = "", checkpoint_path: str | None = None):
        super().__init__(config)

        self.prefix = prefix
        self.quant_config = None  # For future quantization support

        self.embedding_concat_strategy = config.arch_config.embedding_concat_strategy
        self.n_layers_per_group = config.arch_config.n_layers_per_group
        self.num_embedding_padding_tokens = config.arch_config.num_embedding_padding_tokens

        config_path = checkpoint_path if checkpoint_path else config.tokenizer_type

        logger.info("Initializing Reason1TextEncoder (Qwen2.5-VL) from %s", config_path)
        try:
            from transformers import AutoConfig as HFAutoConfig
            hf_config = HFAutoConfig.from_pretrained(
                config_path,
                trust_remote_code=True,
            )
        except Exception as e:
            logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.", config_path, e)
            hf_config = Qwen2_5_VLConfig(
                hidden_size=3584,
                intermediate_size=18944,
                max_window_layers=28,
                num_attention_heads=28,
                num_hidden_layers=28,
                num_key_value_heads=4,
                tie_word_embeddings=False,
                vocab_size=152064,
            )

        hf_config.output_hidden_states = True

        if hasattr(config.arch_config, '_attn_implementation') and config.arch_config._attn_implementation:
            hf_config._attn_implementation = config.arch_config._attn_implementation
        else:
            hf_config._attn_implementation = "flash_attention_2"
        logger.info("Reason1 attention implementation: %s", getattr(hf_config, "_attn_implementation", None))

        with torch.device("meta"):
            self.model = Qwen2_5_VLForConditionalGenerationSimple(hf_config)

        self.processor = AutoProcessor.from_pretrained(
            config_path,
            trust_remote_code=True,
        )

        weights_override = envs.FASTVIDEO_REASON1_WEIGHTS_PATH.get()
        if weights_override:
            self.secondary_weights = (_WeightsSource(
                model_or_path=weights_override,
                prefix="",
                fall_back_to_pt=True,
                allow_patterns_overrides=None,
            ), )
            logger.info("Reason1TextEncoder: overlaying weights from %s", weights_override)

        self._weights_loaded = False

    def forward(
        self,
        input_ids: torch.Tensor | None,
        position_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs,
    ) -> BaseEncoderOutput:
        # Cosmos2.5 alignment: keep attention_mask=None.
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=None,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            output_hidden_states=True,
            return_dict=True,
            pixel_values=kwargs.get('pixel_values', None),
            pixel_values_videos=kwargs.get('pixel_values_videos', None),
            image_grid_thw=kwargs.get('image_grid_thw', None),
            video_grid_thw=kwargs.get('video_grid_thw', None),
        )

        hidden_states = outputs.hidden_states
        last_hidden_state = hidden_states[-1]

        return BaseEncoderOutput(
            last_hidden_state=last_hidden_state,
            hidden_states=hidden_states if output_hidden_states else None,
            attention_mask=None,
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        first_weight = None
        weights_list = []
        for name, weight in weights:
            if first_weight is None:
                first_weight = weight
                self.model = self.model.to_empty(device=weight.device)
                self.model.init_weights(buffer_device=weight.device)
            weights_list.append((name, weight))

        params_dict = dict(self.model.named_parameters())
        loaded_params: set[str] = set()
        skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0}

        for name, loaded_weight in weights_list:
            if "lm_head" in name:
                skipped_weights["lm_head"] += 1
                continue
            if "visual" in name:
                skipped_weights["visual"] += 1
                continue
            if "decoder" in name:
                skipped_weights["decoder"] += 1
                continue

            # Handle stacked params mapping (for quantized models)
            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                if name.endswith(".bias") and name not in params_dict:
                    continue

                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
                loaded_params.add(param_name_with_prefix)
                break
            else:
                if name.endswith(".bias") and name not in params_dict:
                    continue

                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
                loaded_params.add(param_name_with_prefix)
        if first_weight is not None:
            self.model = self.model.to(first_weight.device)

        all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
                         for name in params_dict.keys())
        loaded_params.update(all_params)

        # Mark weights as loaded
        self._weights_loaded = True
        return loaded_params

    def compute_text_embeddings_online(
        self,
        data_batch: dict[str, list[str]],
        input_caption_key: str,
    ) -> torch.Tensor:
        prompts = data_batch[input_caption_key]
        return self.compute_text_embeddings(prompts)

    def compute_text_embeddings(
        self,
        prompts: list[str],
        device: str | torch.device = "cuda",
    ) -> torch.Tensor:
        """Compute embeddings for a list of prompts."""
        input_ids_batch = []

        tok = getattr(self.processor, "tokenizer", None)
        if tok is None:
            raise RuntimeError("Reason1TextEncoder requires processor.tokenizer")
        pad_id = getattr(tok, "pad_id", None)
        if pad_id is None:
            pad_id = getattr(tok, "pad_token_id", None)
        if pad_id is None:
            pad_id = getattr(self.model.config, "pad_token_id", None)
        if pad_id is None:
            pad_id = 0

        for prompt in prompts:
            conversations = [
                {
                    "role":
                    "system",
                    "content": [{
                        "type": "text",
                        "text": "You are a helpful assistant who will provide prompts to an image generator.",
                    }],
                },
                {
                    "role": "user",
                    "content": [{
                        "type": "text",
                        "text": prompt,
                    }],
                },
            ]

            try:
                tokenizer_output = tok.apply_chat_template(
                    conversations,
                    tokenize=True,
                    add_generation_prompt=False,
                    add_vision_id=False,
                )
            except TypeError:
                tokenizer_output = tok.apply_chat_template(
                    conversations,
                    tokenize=True,
                    add_generation_prompt=False,
                )

            input_ids = _normalize_chat_template_ids(tokenizer_output)

            if self.num_embedding_padding_tokens > len(input_ids):
                pad_len = self.num_embedding_padding_tokens - len(input_ids)
                input_ids = input_ids + [pad_id] * pad_len
            else:
                input_ids = input_ids[:self.num_embedding_padding_tokens]

            input_ids = torch.LongTensor(input_ids).to(device=device)
            input_ids_batch.append(input_ids)

        input_ids_batch = torch.stack(input_ids_batch, dim=0)

        # Cosmos2.5 alignment: keep attention_mask=None.
        target_device = input_ids_batch.device
        try:
            embed_device = self.model.model.embed_tokens.weight.device  # type: ignore[attr-defined]
        except Exception:
            embed_device = None
        if embed_device is not None and embed_device != target_device:
            self.model = self.model.to(target_device)

        with torch.no_grad():
            position_ids, _ = get_rope_index(
                self.model.config,
                input_ids_batch,
                image_grid_thw=None,
                video_grid_thw=None,
                second_per_grid_ts=None,
                attention_mask=None,
            )
            position_ids = position_ids.to(target_device)

            outputs = self.model.model(
                input_ids=input_ids_batch,
                position_ids=position_ids,
                attention_mask=None,
                output_hidden_states=True,
                return_dict=True,
                use_cache=False,
            )
            hidden_states = outputs.hidden_states

        normalized_hidden_states = []
        for layer_idx in range(1, len(hidden_states)):
            normalized_state = self._mean_normalize(hidden_states[layer_idx])
            normalized_hidden_states.append(normalized_state)

        if self.embedding_concat_strategy == "full_concat":
            text_embeddings = torch.cat(normalized_hidden_states, dim=-1)
        elif self.embedding_concat_strategy == "mean_pooling":
            text_embeddings = torch.stack(normalized_hidden_states).mean(dim=0)
        elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat":
            pooled_embeddings = []
            for i in range(0, len(normalized_hidden_states), self.n_layers_per_group):
                group = normalized_hidden_states[i:i + self.n_layers_per_group]
                pooled = torch.stack(group).mean(dim=0)
                pooled_embeddings.append(pooled)
            text_embeddings = torch.cat(pooled_embeddings, dim=-1)
        else:
            raise ValueError(f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}")

        return text_embeddings

    @staticmethod
    def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor:
        return (tensor - tensor.mean(dim=-1, keepdim=True)) / (tensor.std(dim=-1, keepdim=True) + 1e-8)

fastvideo.models.encoders.reason1.Reason1TextEncoder.forward

forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs) -> BaseEncoderOutput
Source
def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -> BaseEncoderOutput:
    # Cosmos2.5 alignment: keep attention_mask=None.
    outputs = self.model(
        input_ids=input_ids,
        attention_mask=None,
        position_ids=position_ids,
        inputs_embeds=inputs_embeds,
        output_hidden_states=True,
        return_dict=True,
        pixel_values=kwargs.get('pixel_values', None),
        pixel_values_videos=kwargs.get('pixel_values_videos', None),
        image_grid_thw=kwargs.get('image_grid_thw', None),
        video_grid_thw=kwargs.get('video_grid_thw', None),
    )

    hidden_states = outputs.hidden_states
    last_hidden_state = hidden_states[-1]

    return BaseEncoderOutput(
        last_hidden_state=last_hidden_state,
        hidden_states=hidden_states if output_hidden_states else None,
        attention_mask=None,
    )

fastvideo.models.encoders.reason1.Reason1TextEncoder.load_weights

load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]
Source
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    first_weight = None
    weights_list = []
    for name, weight in weights:
        if first_weight is None:
            first_weight = weight
            self.model = self.model.to_empty(device=weight.device)
            self.model.init_weights(buffer_device=weight.device)
        weights_list.append((name, weight))

    params_dict = dict(self.model.named_parameters())
    loaded_params: set[str] = set()
    skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0}

    for name, loaded_weight in weights_list:
        if "lm_head" in name:
            skipped_weights["lm_head"] += 1
            continue
        if "visual" in name:
            skipped_weights["visual"] += 1
            continue
        if "decoder" in name:
            skipped_weights["decoder"] += 1
            continue

        # Handle stacked params mapping (for quantized models)
        for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)
            if name.endswith(".bias") and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
            loaded_params.add(param_name_with_prefix)
            break
        else:
            if name.endswith(".bias") and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)
            param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
            loaded_params.add(param_name_with_prefix)
    if first_weight is not None:
        self.model = self.model.to(first_weight.device)

    all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
                     for name in params_dict.keys())
    loaded_params.update(all_params)

    # Mark weights as loaded
    self._weights_loaded = True
    return loaded_params

fastvideo.models.encoders.reason1.Reason1TextEncoder.compute_text_embeddings_online

compute_text_embeddings_online(data_batch: dict[str, list[str]], input_caption_key: str) -> torch.Tensor
Source
def compute_text_embeddings_online(
    self,
    data_batch: dict[str, list[str]],
    input_caption_key: str,
) -> torch.Tensor:
    prompts = data_batch[input_caption_key]
    return self.compute_text_embeddings(prompts)

fastvideo.models.encoders.reason1.Reason1TextEncoder.compute_text_embeddings

compute_text_embeddings(prompts: list[str], device: str | torch.device = 'cuda') -> torch.Tensor

Compute embeddings for a list of prompts.

Source
def compute_text_embeddings(
    self,
    prompts: list[str],
    device: str | torch.device = "cuda",
) -> torch.Tensor:
    """Compute embeddings for a list of prompts."""
    input_ids_batch = []

    tok = getattr(self.processor, "tokenizer", None)
    if tok is None:
        raise RuntimeError("Reason1TextEncoder requires processor.tokenizer")
    pad_id = getattr(tok, "pad_id", None)
    if pad_id is None:
        pad_id = getattr(tok, "pad_token_id", None)
    if pad_id is None:
        pad_id = getattr(self.model.config, "pad_token_id", None)
    if pad_id is None:
        pad_id = 0

    for prompt in prompts:
        conversations = [
            {
                "role":
                "system",
                "content": [{
                    "type": "text",
                    "text": "You are a helpful assistant who will provide prompts to an image generator.",
                }],
            },
            {
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": prompt,
                }],
            },
        ]

        try:
            tokenizer_output = tok.apply_chat_template(
                conversations,
                tokenize=True,
                add_generation_prompt=False,
                add_vision_id=False,
            )
        except TypeError:
            tokenizer_output = tok.apply_chat_template(
                conversations,
                tokenize=True,
                add_generation_prompt=False,
            )

        input_ids = _normalize_chat_template_ids(tokenizer_output)

        if self.num_embedding_padding_tokens > len(input_ids):
            pad_len = self.num_embedding_padding_tokens - len(input_ids)
            input_ids = input_ids + [pad_id] * pad_len
        else:
            input_ids = input_ids[:self.num_embedding_padding_tokens]

        input_ids = torch.LongTensor(input_ids).to(device=device)
        input_ids_batch.append(input_ids)

    input_ids_batch = torch.stack(input_ids_batch, dim=0)

    # Cosmos2.5 alignment: keep attention_mask=None.
    target_device = input_ids_batch.device
    try:
        embed_device = self.model.model.embed_tokens.weight.device  # type: ignore[attr-defined]
    except Exception:
        embed_device = None
    if embed_device is not None and embed_device != target_device:
        self.model = self.model.to(target_device)

    with torch.no_grad():
        position_ids, _ = get_rope_index(
            self.model.config,
            input_ids_batch,
            image_grid_thw=None,
            video_grid_thw=None,
            second_per_grid_ts=None,
            attention_mask=None,
        )
        position_ids = position_ids.to(target_device)

        outputs = self.model.model(
            input_ids=input_ids_batch,
            position_ids=position_ids,
            attention_mask=None,
            output_hidden_states=True,
            return_dict=True,
            use_cache=False,
        )
        hidden_states = outputs.hidden_states

    normalized_hidden_states = []
    for layer_idx in range(1, len(hidden_states)):
        normalized_state = self._mean_normalize(hidden_states[layer_idx])
        normalized_hidden_states.append(normalized_state)

    if self.embedding_concat_strategy == "full_concat":
        text_embeddings = torch.cat(normalized_hidden_states, dim=-1)
    elif self.embedding_concat_strategy == "mean_pooling":
        text_embeddings = torch.stack(normalized_hidden_states).mean(dim=0)
    elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat":
        pooled_embeddings = []
        for i in range(0, len(normalized_hidden_states), self.n_layers_per_group):
            group = normalized_hidden_states[i:i + self.n_layers_per_group]
            pooled = torch.stack(group).mean(dim=0)
            pooled_embeddings.append(pooled)
        text_embeddings = torch.cat(pooled_embeddings, dim=-1)
    else:
        raise ValueError(f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}")

    return text_embeddings