How to customize a forward pass

[christopher5106]

Hi,
I realized that when I use transformer.forward = custom_transformer_forward.__get__(transformer, TransformerModel), I can't use enable_model_cpu_offload() any more, probably I broke one hook doing so.
What is the best pratice to customize the forward pass ?
Thanks

[a-r-r-o-w]

Interesting. This is more of a question that's suitable for pytorch, and is probably not an intended design use case. Something like following should work better IMO (have not tested though):

from diffusers import DiffusionPipeline, SomeTransformerModel

class MyTransformer(SomeTransformerModel):
    def forward(self, ...):
        # My custom implementation

transformer = MyTransformer.from_pretrained(...)
pipe = DiffusionPipeline.from_pretrained(..., transformer=transformer)
pipe.enable_model_cpu_offload()

If you just want to modify the inputs before proceeding with the actual existing forward impl, a forward hook would probably be better.

I don't see why transformer.forward = my_custom_forward approach won't work though. It does disrupt hooks, but you could enable cpu offloading after the overwrite is done instead of before (I haven't tested though).

[christopher5106]

I experienced that weights were still on CPU during inference which raised an torch cast exception.
The reason is because I modify the forward pass after the call to pipe.enable_model_cpu_offload() and I have no other choice in my current implementation.
So what I would like to do, is to customize the forward pass once cpu offload has been enabled, that means, set the correct accelerate hooks again. Thanks a lot