There is one test failing as I run in my environment deployment:

FAILED tests/models/test_idefics3.py::TestIdefics3::test_model_output[False] - RuntimeError: Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu (when checking argument in method wrapper_CUDA__index_select)

My environment has GPU, but I haven't set up any device to run the tests; Maybe good idea to confirm at the tests level all is running under the same device?

I will test on GPU. Thank you for looking into this @flaviabeo !

GraniteConfig is imported but never used. Should be removed.

This handles multiple fallback cases for NaViT position buckets detection.
It would be good to add a comment explaining why these fallbacks are necessary.

These values are already ints from the config, unnecessary casting.

I see that we have duplicate logic in the two if branches (lines 101-113 and 114-123). Could be refactored: simplified version or something like this:

# Simplified version for k, v in input_sd.items(): key = k[len("model."):] if k.startswith("model.") else k if key.startswith("connector."): key = key[len("connector."):] if key.startswith("modality_projection.proj."): key = key[len("modality_projection.proj."):] new_sd[f"connector.proj.{key}"] = v elif key.startswith("modality_projection."): key = key[len("modality_projection."):] new_sd[f"connector.proj.{key}"] = v elif key == "modality_projection": new_sd["connector.proj.weight"] = v else: new_sd[f"connector.{key}"] = v

Complex nested hasattr checks. Consider adding a method to base classes instead:

def reset_parameters(self):
    # Let each component handle its own initialization
    if hasattr(self.vision_tower, 'reset_parameters'):
        self.vision_tower.reset_parameters()
    if hasattr(self.text_model, 'reset_parameters'):
        self.text_model.reset_parameters()
    nn.init.xavier_uniform_(self.connector.proj.weight)

The to_pil function handles many edge cases but has duplicated PIL import checks. Consider extracting to module level.

Minor: Type hint for x_in should be Optional[torch.Tensor] for consistency.

Question: What's the purpose of * (1 - 1e-6)? This should be documented.

String comparison 'cuda' should use dev1.type constant or be more defensive.

When is use_navit_position_buckets expected to be a defined attribute on the HF config? I don't see it in the idefics3 config here.

IMO it may also be a good idea to separate build_siglip_vision_params and build_siglip_idefics3_vision_params for clarity as well

Good call. use_navit_position_buckets is not reliably present on HF Idefics3/SmolVLM vision configs, so we now treat it as optional and avoid assuming it exists.

I split the logic into:

• build_siglip_vision_params(...) for generic SigLIP defaults (explicit flag if present, else default False)
• build_siglip_idefics3_vision_params(...) for idefics3/smolvlm-scoped NaViT inference

build_idefics3_params(...) now uses the idefics3-specific builder. I also added unit coverage in tests/models/hf/test_param_builders.py for:

• default False when attr is absent
• smolvlm max_image_size heuristic enabling NaViT
• no heuristic leakage for non-idefics3 parent configs

For the error messages here, can you either log the full shape, or use negative indices? This will get confusing if a batch dim of 1 is squeezed out somewhere since it would log the the wrong dim for the number of patches as is

Updated this to make error messages robust to squeezed/broadcasted shapes and to include full tensor shapes.

Changes made:

• use stable indexing (shape[-2] for patch count) instead of fixed positional assumptions
• include full shape context in errors, e.g. x.shape=... / patch_embeds.shape=...
• include expected/actual patch counts and H/W/scale in the message

This should make debugging much clearer when an upstream batch/image dimension changes.

Should not be a staticmethod and should just use self.scale

Agreed, updated.

pixel_unshuffle is now an instance method and uses self.scale directly (no @staticmethod and no external scale arg). This keeps connector behavior fully owned by the module config.

Can you try to align this with what is done in llava next, which is to create the merged multimodal embeddings in prepare_inputs_for_generation, using kwargs that have optional values for stuff from the HF processor, e.g., pixel values etc?

I imagine this will likely refactored in the future to be more well patterned, but this will allow for better compatibility with the existing generate, e.g., by passing prepare_inputs_for_generation as a preprocessing hook, and also keep things more consistent with projects that may want to consume it like vLLM spyre

Implemented this alignment with LlavaNext/shared generation flow.

Changes made:

• added prepare_inputs_for_generation(iteration, input_ids, kwargs) to build multimodal embeddings on prefill and avoid re-encoding images during cached decode
• switched Idefics3.generate() to use fms.utils.generation.generate(..., prepare_model_inputs_hook=self.prepare_inputs_for_generation)
• threaded optional processor kwargs (pixel_values, images, pixel_attention_mask) through the hook path

This keeps idefics3 compatible with the shared generate path and closer to the existing multimodal pattern used in LlavaNext.

when is this fallback case expected?

Added clarification and narrowed the fallback expectations.

Expected/default layout is full_model.model.vision_model.
Fallback path (full_model.vision_model) is only for HF wrapper/layout variation where the nested model container is absent.

I documented this inline and include model type context in fallback logging/errors so it’s explicit when/why the fallback is taken.

Can you add some logs to some of these fallback cases?

Added logs on fallback paths.

Current behavior:

• log info when using non-default but expected layout fallback paths
• log context (HF model class/type) in fallback messages
• fail loudly (with explicit error) for risky connector extraction cases rather than silently scanning/selecting ambiguous params

This keeps normal/default path quiet while making fallback behavior diagnosable.

Similar to some of the above comments on aligning with llava next for model structure and not having a custom .generate - I don't think that we should have an explicit wrapper for the HF processor (even a lazy one for conditional dependencies), since it may cause confusion with how the model is to be used with transformers

Agreed, removed the explicit wrapper class.

Changes made:

• deleted SmolVLMPreprocessor wrapper class
• kept a thin helper load_smolvlm_processor(...) that returns HF AutoProcessor directly
• kept load_smolvlm_preprocessor(...) only as a deprecated compatibility alias to avoid immediate breakage for older call sites

So the canonical interface is now the HF processor itself, aligned with your comment.

is this degree of flexibility actually needed for idefic3, or it is written to generically common attributes for vision towers in VLMs?

Tightened this to idefics3/smolvlm-relevant layouts rather than broad generic scanning.

Current order:

1. checkpoint.model.vision_model (default)
2. checkpoint.model.vision_tower (fallback, logged)
3. top-level known attrs (vision_model, vision_tower, vision) (fallback, logged)
4. dict-like wrapper keys only if values are actual nn.Module objects

So flexibility is now constrained to known/expected variants, with logs on non-default selections.

A bit apprehensive about adding so much state dict key manipulation here since the serialization utils are so complex already, and other models use regexes for pattern replacement, e.g., here. 🤔 I am not a big fan of the regexes either, but might be better to do things consistently so that we don't see different behaviors for this model, since it will be hard to debug, especially if more adapter steps are added

@kaoutar55 @flaviabeo @gkumbhat any thoughts?

Aligned this with the existing regex-replacement pattern approach for consistency.

I replaced the branchy connector key remap logic with ordered regex replacement rules (including both model.connector.* and connector.* inputs), then gate on connector. outputs. This keeps behavior explicit while matching the style used elsewhere (e.g. llava_next) for easier debugging/maintenance.

Hi! Yes, agreed @alex-jw-brooks - better to standardize the way of doing these replacements. These can cause warnings when the weights keys are not mapped correctly and can be hard to fix a bug on the mappings.

can you separate most of the nai vt handling (except the early return with no patch attn mask, which is the same as the normal case) into a separate function to keep forward more minimal? it'll be easier to look at other models that use siglip without it, e.g., granite vision

Done. I split most NaViT-specific logic out of SiglipVisionEmbeddings.forward into a helper (_navit_position_ids(...)).

forward now keeps the early return for the non-NaViT/no-mask path and only calls the helper for the NaViT bucketed position-id path. This keeps the main forward path minimal for non-NaViT siglip users.

I see some tests failing:

______ test_config_mapping[SiglipModel-google/siglip-so400m-patch14-384] _______

arch = 'SiglipModel', model_name = 'google/siglip-so400m-patch14-384'

    @pytest.mark.parametrize("arch,model_name", MODEL_CONFIG_MAP.items())
    def test_config_mapping(arch: str, model_name: str):
        from fms.models.hf.utils import infer_model_configuration
    
        cfg_filename = get_file_kwargs_filename(model_name)
        kwargs_path = os.path.join(FMS_CONFIGS_DIR, cfg_filename)
        if not os.path.isfile(kwargs_path):
            raise FileNotFoundError(
                f"Model {model_name} has no kwargs file; have you generated it?"
            )
        with open(kwargs_path, "r") as f:
            expected_kwargs = json.load(f)
        actual_kwargs = infer_model_configuration(model_name, download_weights=False)
    
        for k, actual_v in actual_kwargs.items():
>           expected_v = expected_kwargs[k]
E           KeyError: 'use_navit_position_buckets'

tests/utils/test_config_mapping.py:62: KeyError
----------------------------- Captured stderr call -----------------------------

Fetching 1 files:   0%|          | 0/1 [00:00<?, ?it/s]
Fetching 1 files: 100%|██████████| 1/1 [00:00<00:00, 29.36it/s]
_ test_config_mapping[LlavaNextForConditionalGeneration-ibm-granite/granite-vision-3.2-2b] _

arch = 'LlavaNextForConditionalGeneration'
model_name = 'ibm-granite/granite-vision-3.2-2b'

It seems related to this change: use_navit_position_buckets - something is not inferring correctly this model configuration or the args are getting lost in between set up and call, maybe?