Merge pull request #8 from v0xie/hotfix

Hotfix: revert _forward restore
v0xie · Dec 28, 2023 · 124355d · 124355d
2 parents 0b0568e + 924d7fb
commit 124355d
Showing 1 changed file with 0 additions and 5 deletions.
diff --git a/agentsd/patch.py b/agentsd/patch.py
@@ -375,8 +375,6 @@ def apply_patch(
         if isinstance_str(module, "BasicTransformerBlock"):
             module._old_class_= [module.__class__]
             _old_forward = None
-            if hasattr(module, "_forward"):
-                _old_forward = [module._forward]
             make_tome_block_fn = make_diffusers_tome_block if is_diffusers else make_tome_block
             module.__class__ = make_tome_block_fn(module.__class__, _old_forward)
             module._tome_info = diffusion_model._tome_info
@@ -414,9 +412,6 @@ def remove_patch(model: torch.nn.Module):
             module._tome_info["hooks"].clear()
 
         if module.__class__.__name__ == "ToMeBlock":
-            if hasattr(module, "_old_forward"):
-                if module._old_forward is not None:
-                    module._forward = module._old_forward[0]
             if hasattr(module, "_old__class__"):
                 module.__class__ = module._old__class__[0]
             else: