✨ Add EoMT Model ||  🚨 Fix Mask2Former loss calculation
Author: yaswanth19
Pull request: yes
Labels: name: New model

name: Vision

# What does this PR do?
Fixes #37171  and continuation of #37392


Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

Image segmentation model so cc @qubvel @nielsrogge!

@qubvel A rough draft is ready for inference :hugs: . Now i am adding support for training and they use mask_annealing to determine probability for attn masks. Is it required in this HF implementation also (I don't see for any other model) or are we fine with having a fixed prob for attn mask

Hey @yaswanth19, do you mean the mask probability changes during the model training? Would be nice to have it but not sure it can be easily implemented tbh. Maybe we can just add a callback for a Trainer to change it? similar to learning rate callback (I mean not adding it to the Transformers actually, but to the docs/fine-tuning guide)

>do you mean the mask probability changes during the model training?

Yup exactly, and adding a trainer callback seems to be a good idea. I will check the feasibility of implementation and if simple enough then we can implement in the model itself else pivot to trainer callback.

@yaswanth19 Thanks for your great work!

A few thoughts:
	•	Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid. A fixed probability of 0 removes masked attention but does hurt performance. If mask annealing isn’t feasible right now, I’d suggest setting it to 0 for now. If you do implement it, note that intermediate mask predictions will be required, so the _predict method might need to move into EoMTEncoder. Just make sure intermediate masks are skipped at inference to avoid slowing things down.  
	•	Testing: It might be good to add an integration test that verifies both the mask and class outputs for semantic, instance, and panoptic segmentation.  
	•	Weight initialization: The current approach likely misinitializes parameters like query embeddings (e.g. std=0.02). All parameters outside the ViT should follow PyTorch defaults (e.g. nn.Embedding uses std=1.0).  
	•	ViT backbone: Would it be simpler to use timm for the ViT? This allows loading pre-trained weights when training EoMT from scratch, avoids unnecessary code, and avoids unsupported use cases like training from a randomly initialized ViT.  
	•	Loss function: Could we reuse the Mask2Former loss? Unless it conflicts with Transformers guidelines, this might reduce redundancy.

Let me know your thoughts.

Thanks @tommiekerssies  for your initial thoughts.

> Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid.

* I am not sure abut mask annealing implementation compatibility with transformers natively. As I have said above, in the worst case we can set to 0 or use tariner callback if that's feasible. 

 > Testing: It might be good to add an integration test

Yup, will add the complete test suite once I have a implementation ready.

> Weight initialization: The current approach likely misinitializes parameters like query embeddings 

Thanks for bringing this to my attention, I can make some correction to initialization later on when we have the end-to-end code ready. IMO, most of the user don't init from scratch and will either finetune it or just perform inference. But having said that I will look at timm implementation and will init in the same way

> ViT backbone: Would it be simpler to use timm for the ViT? 

Ideally yes :sweat_smile: But that's not the library coding standard (Don't want to introduce a hard dependency on `timm`). Also using timm backbone directly will not be compatible with all other features that HF ecosystem provides IMO. I am actually referring the HF VIT and timm implementation to get the best of both worlds and as to not introduce any bug.

> Loss function: Could we reuse the Mask2Former loss? 

Transformers has one model one file philosophy and because of that I have copied the Mask2Former loss completely here. It can be subjective call with Modular file in the sense we can expose Mask2Former loss and import it here for EoMT (Will require additional changes in mask2former) but that can be discussed during reviews with the core maintainer.



> Thanks @tommiekerssies for your initial thoughts.
> 
> > Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid.
> 
> * I am not sure abut mask annealing implementation compatibility with transformers natively. As I have said above, in the worst case we can set to 0 or use tariner callback if that's feasible.
> 
> > Testing: It might be good to add an integration test
> 
> Yup, will add the complete test suite once I have a implementation ready.
> 
> > Weight initialization: The current approach likely misinitializes parameters like query embeddings
> 
> Thanks for bringing this to my attention, I can make some correction to initialization later on when we have the end-to-end code ready. IMO, most of the user don't init from scratch and will either finetune it or just perform inference. But having said that I will look at timm implementation and will init in the same way
> 
> > ViT backbone: Would it be simpler to use timm for the ViT?
> 
> Ideally yes 😅 But that's not the library coding standard (Don't want to introduce a hard dependency on `timm`). Also using timm backbone directly will not be compatible with all other features that HF ecosystem provides IMO. I am actually referring the HF VIT and timm implementation to get the best of both worlds and as to not introduce any bug.
> 
> > Loss function: Could we reuse the Mask2Former loss?
> 
> Transformers has one model one file philosophy and because of that I have copied the Mask2Former loss completely here. 

Thanks for the clarifications!

Regarding mask annealing, I agree that 0 for now is fine. That means effectively disabling masked attention and mask annealing, which is what the current code already does, so no changes needed on that front.

For weight initialization and the ViT backbone, I understand the constraints around using timm. In that case, I’d just make sure that the non-ViT parameters (query embeddings, mask MLP, upscale blocks, class head) aren’t using any custom initializations and instead follow PyTorch defaults. Should be a quick fix.

Let me know if you’d like me to look at any part in more detail.

Hi @qubvel ,I’m working on refactoring the training logic for EoMT , and I’m running into a design challenge:

In the original single‐class implementation, they call _predict (which uses the class predictor head) on intermediate layer outputs to build attention masks. Because everything lives in one class, this is straightforward.

Refer: https://github.com/tue-mps/eomt/blob/c311b377d3189c976163e4ceb2156d90bb7db88f/models/eomt.py#L130 

In our modular HF version, the encoder (EoMTEncoder) only runs the transformer blocks, and _predict (with mask_head, upscale_block, and class_predictor) lives in EoMTModel or EoMTForUniversalSegmentation. That separation means the encoder loop can’t access _predict , so we can’t reconstruct the original training flow.

I have two solutions in mind, LMK your thoughts on the below approaches and suggest any other better alternative:

1.) Club all the classes from `EoMTEncoder` into `EomtForUniversalSegmentation`, in this ways we can do all the processing in a single forward class.

2.) Move _predict which includes  mask_head, upscale_block into EoMTEncoder and somehow pass the class_head. Pass these modules into the encoder class so that inside its forward loop it can call _predict, build the attention mask, and feed it into the next block. IMO this is a bit dirty and flow is tangled 😅 .

Here the _predict func is the same code which is used in Mask2Former for get mask_logits and class_logits from model output.

@qubvel A ping for your comments on the above issue/challenge so that I can refactor it this weekend 😅 

@yaswanth19 can you send me your email address so we can set up a Slack channel to discuss the integration, rather than having to discuss everything on the Github thread?

Yup sure @NielsRogge , Here is my email id: yaswanthgali8@gmail.com

Nice work @yaswanth19 ! For fixing the initialisation, maybe you could have a look here, as the Mask2Former implementation in HuggingFace suffers from the same bug: https://github.com/huggingface/transformers/issues/35877#issuecomment-2749038225

Hi @tommiekerssies , Sorry for the delay - was occupied with other personal work - I will fix the initialization part 🤗. The model is ready and logits are matching as of now but code still needs some refactoring. The image processing is also ready but right now it doesn't support targets (masks and labels) and regarding that I have a few questions.

Can you explain what is the CLASS_MAPPING and INSTANCE_MAPPING in each xx_dataset.py files. Also it would be great if you could explain how are the mask and class labels extracted specifically the logic of `target_parser` function. It's quite different for each dataset class and I am having hard time on getting a common/general logic intuition to implement.

Another query is we are padding image two times - once in transforms and once again in function `resize_and_pad_imgs_instance_panoptic`. Why is it so 🤔 

> Hi @tommiekerssies , Sorry for the delay - was occupied with other personal work - I will fix the initialization part 🤗. The model is ready and logits are matching as of now but code still needs some refactoring. The image processing is also ready but right now it doesn't support targets (masks and labels) and regarding that I have a few questions.
> 
> Can you explain what is the CLASS_MAPPING and INSTANCE_MAPPING in each xx_dataset.py files. Also it would be great if you could explain how are the mask and class labels extracted specifically the logic of `target_parser` function. It's quite different for each dataset class and I am having hard time on getting a common/general logic intuition to implement.
> 
> Another query is we are padding image two times - once in transforms and once again in function `resize_and_pad_imgs_instance_panoptic`. Why is it so 🤔

Hi @yaswanth19, great work!

**CLASS_MAPPING** is used to remap the dataset’s class IDs to a contiguous range without gaps. **INSTANCE_MAPPING** is specific to ADE20K panoptic, which requires merging semantic and instance labels: we apply `CLASS_MAPPING` to the semantic labels (skipping “thing” classes) and `INSTANCE_MAPPING` to the instance labels (which include “thing” classes).

The `target_parser` function is dataset-specific. It converts labels from the dataset’s original format into the consistent `(masks, labels)` format expected by our training/evaluation code. For example, COCO instance annotations use RLE in JSON, while COCO panoptic labels are stored as PNGs. There are also quirks like COCO’s `is_crowd` flag, which needs special handling (ignored during training, but used for proper AP/PQ evaluation).

Due to these differences, it might not be ideal to include dataset-specific parsing logic in the HF Transformers codebase. Instead, the preprocessor can expect targets to already be in `(masks, labels)` format, leaving the conversion to users. What do you think?

Regarding padding: it’s not applied twice. The transform pipeline is used during training, where we apply random scale jittering, pad to square, and then crop. The `resize_and_pad_imgs_instance_panoptic` function is used only during evaluation, where we resize the long side to the input size and pad the short side to make the image square.

Let me know if anything’s unclear!

Hi @yaswanth19 ! Is this ready for a review? Feel free to ping me when it is!

@qubvel @NielsRogge The PR is ready is review. I have added some Todo's for me which are independent things and not a high priority. I will do them in parallel with reviews once I find some more time.

@tommiekerssies Please have a look at it if you have some bandwidth. IMO focus more the processing class because that's the module which differs from original implementation. AFAIK I have standardized the inference pre-processing correctly and LMK if you find any inconsistencies w.r.t pre and post processing. 

Dear @yaswanth19 , great work! I’m currently on holiday and will review on Tuesday June 3rd when I’m back at work.

The diff between modular and modeling file can be fixed when PR #38503 is merged which will make the CI happy. Fixed the previously failiing tests and all good to go 🚀  @yonigozlan 

Hi @yaswanth19, I haven't finished my full review yet (for which I'd also like to test the code myself). However, one thing which might be useful to already mention, is that the current code doesn't seem to support DINOv2 giant, as the giant model uses a SwiGLU FFN. Should be relatively easy to support this though.

Additionally I think the current weight init code will still cause nn.Embedding classes (such as the queries) to be initialized with 0.02 std, which should be default PyTorch init of 1.0 std.

> support DINOv2 giant, as the giant model uses a SwiGLU FFN

Ohh, I used large variants for testing, hence missed this one. Will make changes to support it, which should be easy.

Thanks @tommiekerssies for the review. Added my replies and my only concern is regarding the post-processor for instance segmentation and attn_mask_prob.

More details: https://github.com/huggingface/transformers/pull/37610#discussion_r2129262453

@yonigozlan The PR is ready for review. 

@yonigozlan Addressed the last comments 🤗 and should be ready for @Cyrilvallez review.

I ran `test_batching_equivalence` multiple times locally but it did not fail for me 🤔 and neither is the case in CI, Seems weird when considering different hardwares. Also, updated the HUB checkpoints with fast image processor.

Also @yonigozlan can you trigger the CI so as to run all the slow tests to check if any other tests are failing or not. I don't have the capacity and env to run all the slow tests 😅 

@Cyrilvallez @yonigozlan Now following camelcasing for model naming and replied to the query. I need to update the ckpts - waiting for your approval so that I can convert all the checkpoints and then @tommiekerssies can transfer the ckpts to `tue-mps` org 🚀 🤗 

@Cyrilvallez The image processor is quite different from anything we have at the moment, as maskformer and mask2former don't have a fast image processor yet. However some modular refactoring should be possible once we get to these!
But for now we wouldn't gain anything with modular so happy to merge as is

@Cyrilvallez A gentle ping for your review/approval so we can complete the final steps of checkpoints transfer and get this merged soon 🤗 🚀 

Hi @tommiekerssies  Can you please add me to the `tue-mps` org so I can transfer the checkpoints from my account to org account. The checkpoints are:

* yaswanthgali/ade20k_panoptic_eomt_large_640
* yaswanthgali/coco_panoptic_eomt_large_640-hf
* yaswanthgali/ade20k_semantic_eomt_large_512-hf
* yaswanthgali/coco_instance_eomt_large_640-hf

I've also added a draft of the README—please feel free to review and suggest edits (Your can do it yourself once the checkpoints are in org). Once the checkpoints are transferred, I’ll update the model_id references accordingly.

Do you want to keep this new checkpoints along with existing ones then these new ones can be suffixed with "-hf" like `tue-mps/coco_panoptic_eomt_large_640-hf` or do you want to replace the existing one?

Additionally, there are a few more checkpoints that still need to be converted to the HF format. If it works for you, it would be great if you could upload those directly under the `tue-mps` org.

Hi @yaswanth19, sure I can give you access. But do we need to have seperate checkpoints for our implementation and Hugging Face? Are the state_dicts actually different, or is it just the config.json and readme.md that you updated?

And one final comment:
https://github.com/huggingface/transformers/pull/37610/files#r2168804369

>Are the state_dicts actually different, 

Yes, the keys would be different. Please refer the key mapping defined below:
https://github.com/yaswanth19/transformers/blob/f0dd8eca910985a79c43ad8b7848681447543e4c/src/transformers/models/eomt/convert_eomt_to_hf.py#L30-L58

>But do we need to have seperate checkpoints for our implementation and Hugging Face

Regarding this - it's your call 😅 

> > Are the state_dicts actually different,
> 
> Yes, the keys would be different. Please refer the key mapping defined below: https://github.com/yaswanth19/transformers/blob/f0dd8eca910985a79c43ad8b7848681447543e4c/src/transformers/models/eomt/convert_eomt_to_hf.py#L30-L58
> 
> > But do we need to have seperate checkpoints for our implementation and Hugging Face
> 
> Regarding this - it's your call 😅

I figured the implementation converted the keys online, thus not requiring to update the checkpoints as they are now. But this is used for offline conversion only?

Yep, The model keys conversion is done **offline**, using the `convert_eomt_to_hf.py` script. It's not on the fly mapping , so we need to convert the checkpoints and store them separately.

Cmd which I am using: `python src/transformers/models/eomt/convert_eomt_to_hf.py --hf_repo_id "" --output_dir ""  --safe_serialization`

> Yep, The model keys conversion is done **offline**, using the `convert_eomt_to_hf.py` script. It's not on the fly mapping , so we need to convert the checkpoints and store them separately.
> 
> Cmd which I am using: `python src/transformers/models/eomt/convert_eomt_to_hf.py --hf_repo_id "" --output_dir "" --safe_serialization`



> Yep, The model keys conversion is done **offline**, using the `convert_eomt_to_hf.py` script. It's not on the fly mapping , so we need to convert the checkpoints and store them separately.
> 
> Cmd which I am using: `python src/transformers/models/eomt/convert_eomt_to_hf.py --hf_repo_id "" --output_dir "" --safe_serialization`

Understood. Is there an easy way to make it online, or is this not the way it should be for Transformers library?

AFAIK you can directly write to hub using `output_hub_path` arg in conversion script. In that case only update would be readme afterwards depending on the checkpoint

> AFAIK you can directly write to hub using `output_hub_path` arg in conversion script. In that case only update would be readme afterwards depending on the checkpoint

What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

>What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

Ahh, AFAIK I don't think that happens. We have to save the HF compatible ckpts in hub manually. CC: @yonigozlan here if there is a better way or can the conversion happen implicitly 🤔 .

Hey @yaswanth19 and @tommiekerssies !

> What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

No that's not how the conversion work in Transformers, the conversion script is only executed manually, and not called by from_pretrained().

I think the best thing to do here, since the current checkpoint doesn't use `model.safetensors`, is simply to upload all the files from @yaswanth19 checkpoints to the current `tue-mps` checkpoints, as there shouldn't be any conflicting files (we should be able to merge the two config.json without issues from what I'm seeing).

That means that when instantiating the model with from_pretrained() without using the `trust_remote_code` argument, you will get the Transformers version of the model (`pytorch_model.bin` won't be downloaded) and you'll get the current `pytorch_model.bin` weights when using it (not sure if trust_remote_code is used in the original repo currently @tommiekerssies ? Or if the checkpoints are just used to manually download the weights). In any case, uploading the Transformers checkpoint to the current checkpoints should work and that way we avoid confusion by having distinct checkpoints.

Thank you @yonigozlan , Indeed we can keep both in the same repo. I was of the opinion that the bin file would be loaded if both are in same repo 😅   or error will pop up but hah `model.safetensors` would be loaded 👌. 

Sounds like a good solution! @yaswanth19 I’ve given you write access to the org. Current configs aren’t used in our code at all, so they can be overwritten with whatever is needed for the Transformers library. So long as the pytorch_model.bin stay there the original code should keep working.

Thanks @tommiekerssies , I have updated those 4 checkpoints which I had. Can you please check everything at your end - that you are able to load the same checkpoint for your implementation. You might need to add `trust_remote_code=True` or only download `model.bin` and load it 🤔 .  Also made the suggested change in model's forward - please have a look. 

> Thanks @tommiekerssies , I have updated those 4 checkpoints which I had. Can you please check everything at your end - that you are able to load the same checkpoint for your implementation. You might need to add `trust_remote_code=True` or only download `model.bin` and load it 🤔 . Also made the suggested change in model's forward - please have a look.

Sorry, I made one final small comment about that change in model's forward ;)
Regarding the checkpoints, I see the model.bin are still there, so that works for our code.
For the remaining checkpoints, would you like me to convert them too using the conversion script in this PR and upload their respective files (safetensors, readme and configs with the right settings)?
Maybe it's a good idea to add a test case for each, so we know they all work? I guess just checking if they can be loaded is enough.

> I made one final small comment about that change in model's forward ;)

Yep, Done @tommiekerssies  🎉 

> Maybe it's a good idea to add a test case for each, so we know they all work? I guess just checking if they can be loaded is enough.

That should be covered in integration tests. I have tested the inference by pulling tue-mps updated repo and it's working fine 🤗 .

> would you like me to convert them too using the conversion script in this PR

yes please, that would be great 😅 . I haven't tried the latest 7B one but assuming it's just a scaled up arc and in that case the conversion script should work fine

@Cyrilvallez  Please do the honours and merge the PR 🚀 🤗 .

@yaswanth19, congrats on the model merged 🎉 🎉 🎉 

Thanks @qubvel :pray: :hugs:  

Congrats 🎉!