Method encode_condition

models/model.py:442–524

(
        self, input_ids, attention_mask, mllm_output, action_cond_features, current_image_index=None, **kwargs
    )

    def encode_condition(
        self, input_ids, attention_mask, mllm_output, action_cond_features, current_image_index=None, **kwargs
    ):
        prompt_embeds = mllm_output.hidden_states
        embeddings = mllm_output.hidden_states[0]

        repeats = 1 if action_cond_features is None else action_cond_features.shape[0] // input_ids.shape[0]
        if current_image_index is None:
            current_image_index = torch.full(
                (input_ids.shape[0],),
                1 if self.config.use_history_obs else 0,
                dtype=torch.long,
                device=input_ids.device,
            )
        elif not isinstance(current_image_index, torch.Tensor):
            current_image_index = torch.tensor(current_image_index, dtype=torch.long, device=input_ids.device)
        else:
            current_image_index = current_image_index.to(device=input_ids.device, dtype=torch.long)

        if repeats > 1:
            input_ids = input_ids.repeat_interleave(repeats, dim=0)
            attention_mask = attention_mask.repeat_interleave(repeats, dim=0)
            embeddings = embeddings.repeat_interleave(repeats, dim=0)
            prompt_embeds =[p.repeat_interleave(repeats, dim=0) for p in prompt_embeds]
            current_image_index = current_image_index.repeat_interleave(repeats, dim=0)

        if self.tokenizer.num_metaqueries > 0:
            # Get positions for all sequences in batch at once
            boi_pos = torch.where(input_ids == self.boi_token_id)[1]
            eoi_pos = torch.where(input_ids == self.eoi_token_id)[1]
            
            def get_vision_positions(input_ids, token_id, image_indices):
                positions = torch.full((input_ids.size(0),), -1, dtype=torch.long, device=input_ids.device)
                rows, cols = torch.where(input_ids == token_id)
                for r in rows.unique():
                    cols_r = cols[rows == r].sort().values
                    image_idx = min(int(image_indices[r].item()), cols_r.numel() - 1)
                    positions[r] = cols_r[image_idx]
                return positions

            vision_start = get_vision_positions(input_ids, self.vision_start_token_id, current_image_index)
            vision_end = get_vision_positions(input_ids, self.vision_end_token_id, current_image_index)


            # Create mask for selecting tokens between BOI and EOI
            batch_size, seq_len = input_ids.shape
            indices = torch.arange(seq_len, device=input_ids.device)[None, :].expand(
                batch_size, -1
            )

            prompt_embeds_mask = (indices > boi_pos[:, None]) & (indices < eoi_pos[:, None])
            embeddings_mask = (indices > vision_start[:, None]) & (indices < vision_end[:, None])

            embeddings = embeddings[embeddings_mask].view(
                batch_size, -1, embeddings.size(-1)
            )
            
            prompt_embeds_all_layer = []