Use Builtin compile_unet of stable-fast

[chengzeyi]

Hey, gameltb.

Thanks for the recent improvements of this plugin.
I have migrated the compile_unet function to be a builtin function of stable-fast.
It is better to use this to avoid hijacking too much, which might cause many backwards compatibility issues.

I tried to rewrite the model offloading logic, but don't know if it really works.
Other parts should work fine except that even if the user remove the loader after compiling,
it will still take effect. To make compiled TorchScript Graphs visible to the outside, I have add an attribute called _traced_modules in the return value of lazy_trace to expose it.

But I don't know how to integrate it with unet_function_wrapper and am still not sure if I do the right things. Maybe there could be other better solutions to avoid hijacking too much.

[gameltb]

Hello, thank you for the PR.

I don't like the idea of changing the diffusion_model permanently, in some cases there are closures that will be passed in as arguments, and changing them permanently will result in an error, even if the user disables the node.

I'm still looking for a suitable way to manage resources and I think it's better for nodes to use function -> ScriptModule and ScriptModule -> CUDAGraph api and let it manage them on a case-by-case basis.

[chengzeyi]

Hello, thank you for the PR.

I don't like the idea of changing the diffusion_model permanently, in some cases there are closures that will be passed in as arguments, and changing them permanently will result in an error, even if the user disables the node.

I'm still looking for a suitable way to manage resources and I think it's better for nodes to use function -> ScriptModule and ScriptModule -> CUDAGraph api and let it manage them on a case-by-case basis.

Could the unet_function be wrapper into a torch.nn.Module? Then compile_unet should work on that.

[gameltb]

You can try it with the FreeU node in ComfyUI.

[chengzeyi]

Now I have optimized the compilation time in the nightly version.
I think in some circumstances it should be OK to trigger recompile, or model offloading logic would be really complex to implement.