fix set_hook and clean_hook

lldacing · Jan 13, 2025 · 42a22bb · 42a22bb
1 parent 0bd11ec
commit 42a22bb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pulidflux.py b/pulidflux.py
@@ -351,15 +351,15 @@ def fix_pulid_patch(self, model):
 def set_hook(diffusion_model, target_forward_orig):
     # comfy.ldm.flux.model.Flux.old_forward_orig_for_pulid = comfy.ldm.flux.model.Flux.forward_orig
     # comfy.ldm.flux.model.Flux.forward_orig = pulid_forward_orig
-    diffusion_model.old_forward_orig_for_pulid = types.MethodType(diffusion_model.forward_orig, diffusion_model)
+    diffusion_model.old_forward_orig_for_pulid = diffusion_model.forward_orig
     diffusion_model.forward_orig = types.MethodType(target_forward_orig, diffusion_model)
 
 def clean_hook(diffusion_model):
     # if hasattr(comfy.ldm.flux.model.Flux, 'old_forward_orig_for_pulid'):
     #     comfy.ldm.flux.model.Flux.forward_orig = comfy.ldm.flux.model.Flux.old_forward_orig_for_pulid
     #     del comfy.ldm.flux.model.Flux.old_forward_orig_for_pulid
     if hasattr(diffusion_model, 'old_forward_orig_for_pulid'):
-        diffusion_model.forward_orig = types.MethodType(diffusion_model.old_forward_orig_for_pulid, diffusion_model)
+        diffusion_model.forward_orig = diffusion_model.old_forward_orig_for_pulid
         del diffusion_model.old_forward_orig_for_pulid
 
 def pulid_outer_sample_wrappers_with_override(wrapper_executor, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):