Skip to content

Refectored modeling_imagegpt.py to enable hooks to capture_outputs#44076

Open
mmahjoub5 wants to merge 1 commit intohuggingface:mainfrom
mmahjoub5:refactor/imagegpt-output-tracing
Open

Refectored modeling_imagegpt.py to enable hooks to capture_outputs#44076
mmahjoub5 wants to merge 1 commit intohuggingface:mainfrom
mmahjoub5:refactor/imagegpt-output-tracing

Conversation

@mmahjoub5
Copy link
Copy Markdown

@mmahjoub5 mmahjoub5 commented Feb 17, 2026

What does this PR do?

This PR refactors the ImageGPT implementation to adopt the new @capture_outputs and @can_return_tuple decorators for standardized output collection, as part of #43979.

The goal is to align Imagegpt with the new output capturing infrastructure and remove legacy return_dict / output_hidden_states handling where possible.


Changes

  1. Integrated ImageGPT with _can_record_outputs

    • Registered stage-level outputs (e.g. hidden_states, attention_weights, 'attention_outputs') using OutputRecorder.
  2. Added @capture_outputs to ImageGPTModel.forward

    • Enables standardized output collection via the new capturing system.
  3. Reduced reliance on legacy output flags

    • Removed most dependencies on return_dict and output_hidden_states across forward functions.
    • The dependency is currently retained in FocalNetEncoder for compatibility and reshaping logic.

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • [x ] Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

@molbap

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

 1. Add _can_record_outputs = {"hidden_states":ImageGPTBlock, "attentions":ImageGPTAttention} on the PreTrainedModel subclass. Use class references, not strings
2. added capture_outputs to main loop
3. added can_return_tuple to higher level fnctions
4. dropped output_attentions, output_hidden_states, return_dict from
   signatures unless required
@github-actions
Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: imagegpt

@mmahjoub5 mmahjoub5 changed the title Refectored modeling_focalnet.py to enable hooks to capture_outputs Refectored modeling_imagegpt.py to enable hooks to capture_outputs Feb 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant