[Development] Is there a way to undo node patches when the run is finished?

[JasonS09]

Hello. I have developed a custom node that patches the attention layer 1 and the unet function (using set_model_unet_function_wrapper and set_model_attn1_patch from the ModelPatcher class). I want this to affect the generation from start to finish of the workflow. However, if I remove that node for the next generation, the patches are still on. The patched functions are running instead of the original functions. I don't want this. I have to reset the server in order to run things as normal.

I'm aware that I could just pop the patches from the ModelPatcher class at the end of the forwarding loop (timestep == 0), but if I do so, patches will be off for the rest of the run (if i want to use another KSampler for example) and I don't want this either. Hope I explained myself well.

[comfyanonymous]

Make sure you do a .clone() like this: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy_extras/nodes_hypernetwork.py#L101 on the ModelPatcher object.