Reusable custom heads, wrappers, and `from_pretrained`?

I have a (complicated) custom head type which I would like to test with multiple base models. I would like to achieve the following:

  1. Easily apply this custom head to different base models
  2. Define custom functions (similar to generate()) on this custom head
  3. Be able to do from_pretrained() to load the base model and randomly initialise the new head
  4. Be able to save and load the complete finetuned model with custom head using the from_pretrained() mechanism

Points 1 and 2 point me towards making a wrapper model.
Points 3 and 4 point me towards inheriting from the base model.

What’s the best practice in this situation? Is there some trick to make the from_pretrained() magic work with the wrapper approach?

One option would be to make a function which automatically creates an inherited model given the base class, but that doesn’t seem like the neat thing to do.