What does this PR do?

Pix2StructTextAttention's weight initialization in _init_weights computes key_value_proj_dim with a fallback that doesn't match the pattern used everywhere else in the same block:

key_value_proj_dim = (
    self.config.text_config.d_kv if isinstance(self.config, Pix2StructConfig) else self.config.hidden_size
)

Every other line in this elif isinstance(module, Pix2StructTextAttention): block mirrors the same leaf attribute name in both branches:

hidden_size = self.config.text_config.hidden_size if isinstance(self.config, Pix2StructConfig) else self.config.hidden_size
n_heads = self.config.text_config.num_heads if isinstance(self.config, Pix2StructConfig) else self.config.num_heads

but key_value_proj_dim falls back to self.config.hidden_size instead of self.config.d_kv, even though Pix2StructTextConfig.d_kv is a distinct field (default 64, versus hidden_size which defaults much higher).

#Impact

When Pix2StructTextModel (or Pix2StructTextAttention) is instantiated standalone — i.e. with a plain Pix2StructTextConfig rather than wrapped in the composite Pix2StructConfig the isinstance check fails, and key_value_proj_dim silently becomes hidden_size instead of d_kv. This gives the wrong initialization std for the query, key, value, and output projection weights.

I verified this numerically: with hidden_size=64, d_kv=16, the observed init std for query.weight matched factor * (hidden_size * hidden_size) ** -0.5 (the buggy formula) rather than factor * (hidden_size * d_kv) ** -0.5 (correct, matching the Pix2StructConfig-wrapped path).

@Cyrilvallez @Rocketknight1