Skip to content

[FEAT] Refactor CUSTOM_ARCHITECTURES to provide custom model#1232

Closed
HarryTheWulf wants to merge 1 commit intohuggingface:mainfrom
HarryTheWulf:add-custom-model-arch
Closed

[FEAT] Refactor CUSTOM_ARCHITECTURES to provide custom model#1232
HarryTheWulf wants to merge 1 commit intohuggingface:mainfrom
HarryTheWulf:add-custom-model-arch

Conversation

@HarryTheWulf
Copy link
Copy Markdown
Contributor

This PR will allow for a custom model to be provided for any of the custom architectures - providing null will default to PreTrainedModel.

This came about when I was creating a new model; I wanted to set it up as an Image Segmentation model.

Below shows how a new architecture can be provided with a specific model, rather than PreTrainedModel.

/** @type {CustomArchitectures} */
const CUSTOM_ARCHITECTURES = new Map([
    ['modnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
    ['birefnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
    ['isnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
    ['u2net', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, U2NetModel]],
    ['ben', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
]);

Providing nothing will default to PreTrainedModel.
Comment thread src/pipelines.js
* @param {T} task The task defining which pipeline will be returned. Currently accepted tasks are:
* - `"audio-classification"`: will return a `AudioClassificationPipeline`.
* - `"automatic-speech-recognition"`: will return a `AutomaticSpeechRecognitionPipeline`.
* - `"background-removal"`: will return a `BackgroundRemovalPipeline`.
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed that we forgot to add the new pipeline to the docblock.

Comment thread src/models.js
Comment on lines +7920 to 7936
/** @type {CustomArchitectures} */
const CUSTOM_ARCHITECTURES = new Map([
['modnet', MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES],
['birefnet', MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES],
['isnet', MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES],
['ben', MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES],
['modnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
['birefnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
['isnet', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
['ben', [MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, null]],
]);
for (const [name, mapping] of CUSTOM_ARCHITECTURES.entries()) {
mapping.set(name, ['PreTrainedModel', PreTrainedModel])

for (let [name, [mapping, model]] of CUSTOM_ARCHITECTURES.entries()) {
if (!model) {
model = PreTrainedModel;
}
mapping.set(name, [model.constructor.name, model]);
MODEL_TYPE_MAPPING.set(name, MODEL_TYPES.EncoderOnly);
MODEL_CLASS_TO_NAME_MAPPING.set(PreTrainedModel, name);
MODEL_NAME_TO_CLASS_MAPPING.set(name, PreTrainedModel);
MODEL_CLASS_TO_NAME_MAPPING.set(model, name);
MODEL_NAME_TO_CLASS_MAPPING.set(name, model);
}
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This refactor allows us to specify a custom model for any given architecture.
Currently I haven't changed anything, so they will all default to PreTrainedModel.

@HarryTheWulf HarryTheWulf changed the title [FEAT] Update custom architectures to provide custom model [FEAT] Refactor CUSTOM_ARCHITECTURES to provide custom model Mar 13, 2025
@xenova
Copy link
Copy Markdown
Collaborator

xenova commented Mar 18, 2026

We've significantly improve how architectures are registered in Transformers.js v4 🙏
Good spot on the background-removal pipeline name missing from the JSDoc. I've moved that commit over to #1595 (with you credited as co-author)

@xenova xenova closed this Mar 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.

2 participants