# Copyright 2026 the HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Testing suite for the PyTorch MuseGlimmer model.""" import unittest from transformers import ( AutoProcessor, MuseGlimmerConfig, MuseGlimmerForConditionalGeneration, MuseGlimmerModel, is_torch_available, ) from transformers.models.muse_glimmer.configuration_muse_glimmer import MuseGlimmerTextConfig, MuseGlimmerVisionConfig from transformers.testing_utils import ( require_torch, require_torch_accelerator, slow, torch_device, ) from ...test_image_processing_common import load_coco_image from ...test_memory_cleanup_mixin import MemoryCleanupMixin from ...test_modeling_common import floats_tensor from ...vlm_tester import VLMModelTest, VLMModelTester if is_torch_available(): import torch class MuseGlimmerVision2TextModelTester(VLMModelTester): base_model_class = MuseGlimmerModel config_class = MuseGlimmerConfig text_config_class = MuseGlimmerTextConfig vision_config_class = MuseGlimmerVisionConfig conditional_generation_class = MuseGlimmerForConditionalGeneration def __init__(self, parent, **kwargs): kwargs.setdefault("image_token_id", 3) kwargs.setdefault("video_token_id", 4) kwargs.setdefault("num_image_tokens", 1) kwargs.setdefault("patch_size", 2) kwargs.setdefault("patch_temporal", 2) kwargs.setdefault("merge_size", 1) kwargs.setdefault("layer_types", ["full_attention", "sliding_attention"]) kwargs.setdefault("pos_emb_height", 4) kwargs.setdefault("pos_emb_width", 4) kwargs.setdefault("intermediate_size", 37) kwargs.setdefault("projector_hidden_size", 32) super().__init__(parent, **kwargs) self.image_grid_thw = (1, 1, 1) self.out_hidden_size = self.hidden_size * self.merge_size**2 @property def _special_token_ids(self): return super()._special_token_ids | {self.video_token_id} def get_vision_config(self): # `layer_types` is shared with the text config by name, but the vision tower uses # "window_attention" instead of "sliding_attention". config = super().get_vision_config() config.layer_types = ["window_attention"] * (config.num_hidden_layers - 1) + ["full_attention"] return config def create_pixel_values(self, batch_size: int | None = None): # Override to 5D for patch-based models batch_size = batch_size if batch_size is not None else self.batch_size grid_t, grid_h, grid_w = self.image_grid_thw num_patches = batch_size * grid_t * grid_h * grid_w return floats_tensor([num_patches, self.patch_temporal * self.num_channels * self.patch_size**2]) def get_additional_inputs(self, config, input_ids, pixel_values, batch_size: int | None = None): batch_size = batch_size if batch_size is not None else self.batch_size return {"image_grid_thw": torch.tensor([list(self.image_grid_thw)] * batch_size, device=torch_device)} @require_torch class MuseGlimmerVision2TextModelTest(VLMModelTest, unittest.TestCase): model_tester_class = MuseGlimmerVision2TextModelTester def test_reverse_loading_mapping(self): # The vendor checkpoint layout is defined relative to the `model.` prefix, which the base # MuseGlimmerModel serializes without. super().test_reverse_loading_mapping(skip_base_model=True) # `meta-models/Muse-Glimmer-30B` is 29.8B parameters -- 55.5 GiB of bfloat16 weights, so it does not fit on the # single 24 GiB accelerator of the daily CI runner. `device_map="auto"` is what makes the test runnable there: # it fills the accelerator(s) and offloads the remainder to CPU RAM, planned against the `CI_CPU_MEMORY_LIMIT_GB` # budget `conftest.py` caps `psutil` to (60 GiB per accelerator). Loading with an explicit single-device map # instead asks for all 55.5 GiB on one card and raises `torch.OutOfMemoryError` while materializing weights. @slow @require_torch_accelerator class MuseGlimmerIntegrationTest(MemoryCleanupMixin, unittest.TestCase): EXPECTED_TEXT_PREFIX = " to find your gift. The purpose of life is to give it away." EXPECTED_IMAGE_PREFIX = " two cats sleeping on a pink" model_id = "meta-models/Muse-Glimmer-30B" @classmethod def setUpClass(cls): super().setUpClass() cls.model = None @classmethod def get_model(cls): # Materialized on first use and shared by the whole class. Loading per test method paid for 55.5 GiB # of weights twice, and left the first copy alive while the second was being materialized. if cls.model is None: cls.model = MuseGlimmerForConditionalGeneration.from_pretrained( cls.model_id, dtype=torch.bfloat16, device_map="auto" ) return cls.model def setUp(self): super().setUp() self.processor = AutoProcessor.from_pretrained(self.model_id) def test_text_generation_matches_reference(self): # The reference implementation tokenizes raw completions as [bos] + encode(prompt). model = self.get_model() tokenizer = self.processor.tokenizer prompt = "The meaning of life is" prompt_ids = tokenizer(prompt, add_special_tokens=False).input_ids input_ids = torch.tensor([[tokenizer.bos_token_id] + prompt_ids], device=torch_device) # The assertion only compares the first 15 tokens' worth of text, so anything generated past that is # never checked. 24 leaves margin for a different token split without paying for tokens nobody reads -- # which is not free: on a single-accelerator runner every decoded token restreams the CPU-offloaded # weights across PCIe. output = model.generate(input_ids=input_ids, max_new_tokens=24, do_sample=False) completion = tokenizer.decode(output[0, input_ids.shape[1] :], skip_special_tokens=True) self.assertEqual(completion[: len(self.EXPECTED_TEXT_PREFIX)], self.EXPECTED_TEXT_PREFIX) def test_image_generation_matches_reference(self): model = self.get_model() processor = self.processor tokenizer = processor.tokenizer image = load_coco_image("000000039769.jpg").convert("RGB") image_inputs = processor(text="<|patch|>In this photo we can see", images=[image], return_tensors="pt") image_grid_thw = image_inputs["image_grid_thw"] self.assertEqual(image_grid_thw.tolist(), [[1, 34, 46]]) num_vision_tokens = int(image_grid_thw.prod(dim=-1).sum() // processor.image_processor.merge_size**2) self.assertEqual(num_vision_tokens, 391) input_ids = image_inputs["input_ids"].to(torch_device) self.assertEqual((input_ids == processor.image_start_token_id).sum().item(), 1) self.assertEqual((input_ids == model.config.image_token_id).sum().item(), num_vision_tokens) self.assertEqual((input_ids == processor.image_end_token_id).sum().item(), 1) output = model.generate( input_ids=input_ids, pixel_values=image_inputs["pixel_values"].to(torch_device, torch.bfloat16), image_grid_thw=image_grid_thw.to(torch_device), # Same reasoning as the text test: the asserted prefix is 6 tokens, so 12 is margin, and the 48 # further tokens this used to generate were never compared against anything. max_new_tokens=12, do_sample=False, ) completion = tokenizer.decode(output[0, input_ids.shape[1] :], skip_special_tokens=True) self.assertEqual(completion[: len(self.EXPECTED_IMAGE_PREFIX)], self.EXPECTED_IMAGE_PREFIX)