fix(network): reset_parameters checks nonexistent (+3 more)

Author: andrewwhitecdw

fastgen/networks/cosmos_predict2/network.py

@@ -431,6 +431,11 @@ def forward(
             - [output, features]: If feature_indices is non-empty
             - (output, logvar): If return_logvar=True
         """
+        if feature_indices is None:
+            feature_indices = set()
+        if return_features_early and len(feature_indices) == 0:
+            return []
+
         x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb = self.prepare_embedded_sequence(
             x_B_C_T_H_W,
             fps=fps,
@@ -641,9 +646,16 @@ def encode(self, text: Union[str, List[str]], precision: torch.dtype = torch.bfl
             input_ids_batch.append(formatted)
 
         input_ids = torch.cat(input_ids_batch, dim=0).to(self.device)
+        pad_token_id = self.tokenizer.pad_token_id or 0
+        attention_mask = (input_ids != pad_token_id).long()
 
         # Forward pass with hidden states
-        outputs = self.model(input_ids=input_ids, output_hidden_states=True, return_dict=True)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
         hidden_states = outputs.hidden_states  # Tuple of (num_layers + 1) tensors
 
         # Skip embedding layer (index 0), normalize and combine layers
@@ -1007,8 +1019,8 @@ def reset_parameters(self):
                         nn.init.zeros_(m.bias)
 
         # Reinitialize RoPE buffers (non-persistent buffers must be recomputed)
-        if hasattr(self.transformer, "rope_embedder") and self.transformer.rope_embedder is not None:
-            rope = self.transformer.rope_embedder
+        if hasattr(self.transformer, "pos_embedder") and self.transformer.pos_embedder is not None:
+            rope = self.transformer.pos_embedder
             device = next(rope.buffers()).device
 
             # Recompute the sequence buffer
@@ -1327,7 +1339,7 @@ def forward(
             - ([output, features], logvar): If both features and logvar requested
         """
         if feature_indices is None:
-            feature_indices = {}
+            feature_indices = set()
         if return_features_early and len(feature_indices) == 0:
             # Exit immediately if user requested this.
             return []