huggingface
diff --git a/‎src/transformers/models/auto/modeling_auto.py
Lines changed: 2 additions & 1 deletion b/‎src/transformers/models/auto/modeling_auto.py
Lines changed: 2 additions & 1 deletion
diff --git a/‎src/transformers/models/gemma3p5/configuration_gemma3p5.py
Lines changed: 50 additions & 9 deletions b/‎src/transformers/models/gemma3p5/configuration_gemma3p5.py
Lines changed: 50 additions & 9 deletions
diff --git a/‎src/transformers/models/gemma3p5/convert_gemma3p5_weights.py
Lines changed: 164 additions & 16 deletions b/‎src/transformers/models/gemma3p5/convert_gemma3p5_weights.py
Lines changed: 164 additions & 16 deletions
@@ -570,8 +570,9 @@
         ("gemma", "GemmaForCausalLM"),
         ("gemma2", "Gemma2ForCausalLM"),
         ("gemma3", "Gemma3ForConditionalGeneration"),
-        ("gemma3p5", "Gemma3p5ForCausalLM"),
         ("gemma3_text", "Gemma3ForCausalLM"),
+        ("gemma3p5", "Gemma3p5ForConditionalGeneration"),
+        ("gemma3p5_text", "Gemma3p5ForCausalLM"),
         ("git", "GitForCausalLM"),
         ("glm", "GlmForCausalLM"),
         ("glm4", "Glm4ForCausalLM"),
 
@@ -268,20 +268,53 @@ def __init__(
 
 
 class Gemma3p5AudioConfig(PretrainedConfig):
-    model_type = "gemma3p5"
+    model_type = "gemma3p5_audio"
 
     def __init__(
         self,
-        *args,
+        input_feat_size: int = 80,
         hidden_size: int = 1536,
         embedding_norm_eps: float = 1e-6,
-        vocab_size: int = 256_128,
+        vocab_size: int = 128,
+        gradient_clipping: float = 10_000_000_000.0,
+        conf_attention_chunk_size: int = 12,
+        conf_attention_context_left: int = 13,
+        conf_attention_context_right: int = 0,
+        conf_attention_invalid_logits_value: float = -1.0e9,
+        conf_attention_logit_cap: float = 50.0,
+        conf_num_attention_heads: int = 8,
+        conf_num_hidden_layers: int = 12,
+        conf_conv_kernel_size: int = 5,
+        conf_positional_bias_size: int = 256,
+        conf_reduction_factor: int = 4,
+        conf_residual_weight: float = 0.5,
+        sscp_conv_channel_size: tuple[int, int] = (128, 32),
+        sscp_conv_group_norm_eps: float = 1e-3,
+        sscp_conv_kernel_size: tuple[tuple[int, int], tuple[int, int]] = ((3, 3), (3, 3)),
+        sscp_conv_stride_size: tuple[tuple[int, int], tuple[int, int]] = ((2, 2), (2, 2)),
         **kwargs,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(**kwargs)
+        self.input_feat_size = input_feat_size
         self.hidden_size = hidden_size
         self.embedding_norm_eps = embedding_norm_eps
         self.vocab_size = vocab_size
+        self.gradient_clipping = gradient_clipping
+        self.conf_attention_chunk_size = conf_attention_chunk_size
+        self.conf_attention_context_left = conf_attention_context_left
+        self.conf_attention_context_right = conf_attention_context_right
+        self.conf_attention_invalid_logits_value = conf_attention_invalid_logits_value
+        self.conf_attention_logit_cap = conf_attention_logit_cap
+        self.conf_num_attention_heads = conf_num_attention_heads
+        self.conf_num_hidden_layers = conf_num_hidden_layers
+        self.conf_conv_kernel_size = conf_conv_kernel_size
+        self.conf_positional_bias_size = conf_positional_bias_size
+        self.conf_reduction_factor = conf_reduction_factor
+        self.conf_residual_weight = conf_residual_weight
+        self.sscp_conv_channel_size = sscp_conv_channel_size
+        self.sscp_conv_eps = sscp_conv_group_norm_eps
+        self.sscp_conv_kernel_size = sscp_conv_kernel_size
+        self.sscp_conv_stride_size = sscp_conv_stride_size
 
 
 class Gemma3p5VisionConfig(PretrainedConfig):
@@ -369,8 +402,11 @@ def __init__(
         audio_soft_tokens_per_image: int = 256,
         vision_soft_tokens_per_image: int = 256,
         boi_token_id: int = 255_999,
-        eoi_token_id: int = 256_000,
-        image_token_id: int = 262_144,
+        eoi_token_id: int = 262_144,
+        image_token_id: int = 262_145,
+        boa_token_id: int = 256_000,
+        eoa_token_id: int = 262_272,
+        audio_token_id: int = 262_273,
         initializer_range: float = 0.02,
         **kwargs,
     ):
@@ -385,12 +421,14 @@ def __init__(
         if isinstance(vision_config, dict):
             vision_config = Gemma3p5VisionConfig(**vision_config)
         elif vision_config is None:
-            logger.info("vision_config is None. Vision capabilities will not be used.")
+            vision_config = Gemma3p5VisionConfig()
+            logger.info("vision_config is None. Using default Gemma3p5VisionConfig.")
 
         if isinstance(audio_config, dict):
             audio_config = Gemma3p5AudioConfig(**audio_config)
         elif audio_config is None:
-            logger.info("audio_config is None. Audio capabilities will not be used.")
+            audio_config = Gemma3p5AudioConfig()
+            logger.info("audio_config is None. Using default Gemma3p5AudioConfig.")
 
         self.text_config = text_config
         self.vision_config = vision_config
@@ -401,7 +439,10 @@ def __init__(
         self.boi_token_id = boi_token_id
         self.eoi_token_id = eoi_token_id
         self.image_token_id = image_token_id
+        self.boa_token_id = boa_token_id
+        self.eoa_token_id = eoa_token_id
+        self.audio_token_id = audio_token_id
         self.initializer_range = initializer_range
 
 
-__all__ = ["Gemma3p5Config", "Gemma3p5TextConfig"]
+__all__ = ["Gemma3p5Config", "Gemma3p5AudioConfig", "Gemma3p5TextConfig", "Gemma3p5VisionConfig"]
@@ -18,9 +18,9 @@
 
 python src/transformers/models/gemma3p5/convert_gemma3p5_weights.py \
     --variant='gemma3p5_4b' \
-    --tokenizer_path="$HOME/nano3/checkpoints/tokenizer/gemma3n_cleaned_262144.spiece" \
-    --checkpoint_path="$HOME/nano3/checkpoints/4b_pt_orbax/" \
-    --output_path="$HOME/nano3/checkpoints/4b_pt_safetensors/"
+    --tokenizer_path="$HOME/gemma3p5/checkpoints/tokenizer/gemma3p5-tokenizer.model" \
+    --checkpoint_path="$HOME/gemma3p5/checkpoints/4b_pt_orbax/" \
+    --output_path="$HOME/gemma3p5/checkpoints/4b_pt_safetensors/"
 """
 
 from collections.abc import Iterator, Sequence
@@ -39,10 +39,11 @@
     Gemma3p5ForConditionalGeneration,
     Gemma3ImageProcessor,
     Gemma3Processor,
+    Gemma3p5AudioConfig,
     Gemma3p5TextConfig,
+    Gemma3p5VisionConfig,
     GemmaTokenizerFast,
     GenerationConfig,
-    SiglipVisionConfig,
 )
 from transformers.image_utils import PILImageResampling
 
@@ -94,6 +95,10 @@
 
 _DTYPES = {"float32", "bfloat16", "float16"}
 
+_AUDIO_ENCODER_PARAMETER = "AudioEncoder/encoder"
+_AUDIO_ENCODER_CONFORMER = f"{_AUDIO_ENCODER_PARAMETER}/conformer/stacked_layers"
+_AUDIO_ENCODER_SSCP = f"{_AUDIO_ENCODER_PARAMETER}/feature"
+
 _TRANSFORMER_PARAMETER = "transformer"
 _TRANSFORMER_ALTUP_PROJ = f"{_TRANSFORMER_PARAMETER}/altup_projection_"
 _TRANSFORMER_ALTUP_UNEMB = f"{_TRANSFORMER_PARAMETER}/altup_unembed_projection_"
@@ -104,10 +109,6 @@
 _TRANSFORMER_POST_TRAINING_PREFIX = "rlx_networks/policy_network/"
 _TRANSFORMER_POST_TRAINING_PREFIX_LEN = len(_TRANSFORMER_POST_TRAINING_PREFIX)
 
-# TODO: ryanmullins - Figure out the vision config
-_VISION_CONFIG = {}
-
-
 _VARIANT_GEMMA_3_2B = "gemma3p5_2b"
 _VARIANT_GEMMA_3_4B = "gemma3p5_4b"
 _VARIANTS = {
@@ -127,16 +128,25 @@
             query_pre_attn_scalar=256,
             max_position_embeddings=32_768,
         ),
-        vision_config=_VISION_CONFIG,
+        vision_config=Gemma3p5VisionConfig(),
+        audio_config=Gemma3p5AudioConfig(),
     ),
     _VARIANT_GEMMA_3_4B: Gemma3p5Config(
         text_config=Gemma3p5TextConfig(),
-        vision_config=_VISION_CONFIG,
+        vision_config=Gemma3p5VisionConfig(),
+        audio_config=Gemma3p5AudioConfig(),
     ),
 }
 
 # ==== Flags ====
 
+_AUDIO_DTYPE = flags.DEFINE_enum(
+    name="audio_dtype",
+    default="bfloat16",
+    help="The floating point precision (aka dtype) of the model.",
+    enum_values=_DTYPES,
+)
+
 _CHECKPOINT_PATH = flags.DEFINE_string(
     name="checkpoint_path",
     default=None,
@@ -190,6 +200,125 @@
 )
 
 
+def convert_audio_encoder_weights(
+    config: Gemma3p5AudioConfig,
+    path: str,
+    param: str,
+    weights: np.ndarray,
+) -> Iterator[tuple[str, np.ndarray]]:
+
+    converted_paths: list[str] = []
+    converted_weights: list[Any] = []
+
+    if path.startswith(_AUDIO_ENCODER_CONFORMER):
+        assert weights.shape[0] == config.conf_num_hidden_layers
+
+        for i, matrix in enumerate(weights):
+            if "fflayer_end" in path:
+                base = f"audio_tower.conformer.{i}.ffw_layer_end"
+
+                if path.endswith("ffn_layer1"):
+                    converted_paths.append(f"{base}.ffw_layer_1.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("ffn_layer2"):
+                    converted_paths.append(f"{base}.ffw_layer_2.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("post_layer_norm"):
+                    converted_paths.append(f"{base}.post_layer_norm.weight")
+                    converted_weights.append(matrix)
+                elif path.endswith("pre_layer_norm"):
+                    converted_paths.append(f"{base}.pre_layer_norm.weight")
+                    converted_weights.append(matrix)
+            elif "fflayer_start" in path:
+                base = f"audio_tower.conformer.{i}.ffw_layer_start"
+
+                if path.endswith("ffn_layer1"):
+                    converted_paths.append(f"{base}.ffw_layer_1.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("ffn_layer2"):
+                    converted_paths.append(f"{base}.ffw_layer_2.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("post_layer_norm"):
+                    converted_paths.append(f"{base}.post_layer_norm.weight")
+                    converted_weights.append(matrix)
+                elif path.endswith("pre_layer_norm"):
+                    converted_paths.append(f"{base}.pre_layer_norm.weight")
+                    converted_weights.append(matrix)
+            elif path.endswith("final_ln"):
+                converted_paths.append(f"audio_tower.conformer.{i}.norm.weight")
+                converted_weights.append(matrix)
+            elif "lconv" in path:
+                base = f"audio_tower.conformer.{i}.lconv1d"
+
+                if path.endswith("conv_norm"):
+                    converted_paths.append(f"{base}.conv_norm.weight")
+                    converted_weights.append(matrix)
+                elif path.endswith("depthwise_conv1d"):
+                    converted_paths.append(f"{base}.depthwise_conv1d.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("linear_end"):
+                    converted_paths.append(f"{base}.linear_end.weight")
+                    converted_weights.append(matrix)
+                elif path.endswith("linear_start"):
+                    converted_paths.append(f"{base}.linear_start.weight")
+                    converted_weights.append(matrix.transpose())
+                elif path.endswith("ln"):
+                    converted_paths.append(f"{base}.pre_layer_norm.weight")
+                    converted_weights.append(matrix)
+            elif "trans_atten" in path:
+                base = f"audio_tower.conformer.{i}.attention"
+
+                if param == "per_dim_scale":
+                    converted_paths.append(f"{base}.attn.per_dim_scale")
+                    converted_weights.append(matrix)
+
+                if path.endswith("query_key_value_projection"):
+                    converted_paths.extend([
+                        f"{base}.attn.q_proj.weight", f"{base}.attn.k_proj.weight", f"{base}.attn.v_proj.weight"
+                    ])
+                    converted_weights.extend([
+                        m.squeeze().reshape(config.hidden_size, config.hidden_size).transpose()
+                        for m in np.split(matrix, 3, axis=1)
+                    ])
+                elif path.endswith("pos_proj"):
+                    converted_paths.append(f"{base}.attn.relative_position_embedding.pos_proj.weight")
+                    converted_weights.append(matrix.reshape(config.hidden_size, config.hidden_size).transpose())
+                elif path.endswith("post"):
+                    converted_paths.append(f"{base}.post.weight")
+                    converted_weights.append(matrix.reshape(config.hidden_size, config.hidden_size).transpose())
+                elif path.endswith("post_norm"):
+                    converted_paths.append(f"{base}.post_norm.weight")
+                    converted_weights.append(matrix)
+                elif path.endswith("pre_norm"):
+                    converted_paths.append(f"{base}.pre_attn_norm.weight")
+                    converted_weights.append(matrix)
+    elif path.startswith(_AUDIO_ENCODER_SSCP):
+        if path.endswith("input_proj"):
+            converted_paths.append(f"audio_tower.subsample_conv_projection.input_proj_linear.weight")
+            converted_weights.append(
+                weights.reshape(config.sscp_conv_channel_size[1] ** 2, config.hidden_size).transpose()
+            )
+        elif "norm_" in path:
+            index = int(path[-1])
+            converted_paths.extend([
+                f"audio_tower.subsample_conv_projection.conv_{index}.norm.bias",
+                f"audio_tower.subsample_conv_projection.conv_{index}.norm.weight",
+            ])
+            converted_weights.extend([np.zeros_like(weights), weights])
+        elif "subsampling_" in path:
+            index = int(path[-1])
+            converted_paths.append(f"audio_tower.subsample_conv_projection.conv_{index}.conv.weight")
+            converted_weights.append(weights.transpose())
+
+    if (cpl := len(converted_paths)) != (cwl := len(converted_weights)):
+        raise ValueError(
+            "The `converted_paths` and `converted_weights` should be the same "
+            f"length. Got {cpl} and {cwl}, respectively, for {path}."
+        )
+
+    return zip(converted_paths, converted_weights)
+
+
 def convert_transformer_weights(
     config: Gemma3p5TextConfig,
     path: str,
@@ -215,7 +344,6 @@ def convert_transformer_weights(
         attention_type_index = int(path[_TRANSFORMER_DECODER_BLOCK_LEN])
         assert weights.shape[0] == config.num_hidden_layers / config.sliding_window_pattern
 
-
         for i, matrix in enumerate(weights):
             layer_idx = config.sliding_window_pattern * i + attention_type_index
             base_path = f"model.layers.{layer_idx}"
@@ -302,7 +430,6 @@ def convert_transformer_weights(
         if param == "input_embedding":
             converted_paths.append("model.embed_tokens.weight")
             converted_weights.append(weights)
-        # TODO: ryanmullins - support multimodal embedding matrices
         elif param == "per_layer_embeddings":
             converted_paths.append("model.embed_tokens_per_layer.weight")
             converted_weights.append(weights.reshape(
@@ -348,11 +475,31 @@ def update_tree(path: str, weights: np.ndarray, target_dtype: torch.dtype) -> No
             )
 
     for (path, param), value in tree.flatten_with_path(ckpt):
-        if path.startswith(_TRANSFORMER_PARAMETER):
+        if param == "audio_input_embedding_extra":
+            update_tree("embed_audio.embedding.weight", value, config.audio_config.torch_dtype)
+        elif path.endswith("audio_embedding_norm"):
+            update_tree("embed_audio.hard_embedding_norm.weight", value, config.audio_config.torch_dtype)
+        elif path.endswith("audio_input_projection"):
+            update_tree("embed_audio.embedding_projection.weight", value.transpose(), config.audio_config.torch_dtype)
+        elif path.endswith("audio_soft_embedding_norm"):
+            update_tree("embed_audio.soft_embedding_norm.weight", value, config.audio_config.torch_dtype)
+        elif param == "mm_input_embedding_extra":
+            update_tree("embed_vision.embedding.weight", value, config.vision_config.torch_dtype)
+        elif path.endswith("mm_hard_embedding_norm"):
+            update_tree("embed_vision.embedding_norm.weight", value, config.vision_config.torch_dtype)
+        elif path.endswith("mm_input_projection"):
+            update_tree(
+                "embed_vision.embedding_projection.weight", value.transpose(), config.vision_config.torch_dtype
+            )
+        elif path.startswith(_TRANSFORMER_PARAMETER):
             for path, weights in convert_transformer_weights(config.text_config, path, param, value):
-                update_tree(path, weights, config.text_config.torch_dtype)
+                update_tree(f"language_model.{path}", weights, config.text_config.torch_dtype)
+        elif path.startswith(_AUDIO_ENCODER_PARAMETER):
+            for path, weights in convert_audio_encoder_weights(config.audio_config, path, param, value):
+                update_tree(path, weights, config.audio_config.torch_dtype)
+
 
-    hf_tree["lm_head.weight"] = hf_tree["model.embed_tokens.weight"]
+    hf_tree["language_model.lm_head.weight"] = hf_tree["language_model.model.embed_tokens.weight"]
 
     return hf_tree
 
@@ -364,6 +511,7 @@ def main(*args):
     variant = _VARIANT.value
 
     config = _VARIANTS[variant]
+    config.audio_config.torch_dtype = getattr(torch, _AUDIO_DTYPE.value)
     config.text_config.torch_dtype = getattr(torch, _TRANSFORMER_DTYPE.value)
     config.vision_config.torch_dtype = getattr(torch, _VISION_DTYPE.value)
     if _INCLUDE_CHAT_TEMPLATE.value:
@@ -381,7 +529,7 @@ def main(*args):
     logging.info("Converted Gemma 3 (%s) state tree from Orbax to Hugging Face.", variant)
 
     with accelerate.init_empty_weights():
-        model = Gemma3p5ForCausalLM(config=config.text_config)
+        model = Gemma3p5ForConditionalGeneration(config=config)
 
     model.load_state_dict(state_tree, assign=True, strict=True)
     logging.info(