Fix ControlNetUnit copy issue (#2910)

* Fix ControlNetUnit copy issue * Add copy test
Mikubill · May 20, 2024 · b197f7c · b197f7c
1 parent a5c0da5
commit b197f7c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/internal_controlnet/args.py b/internal_controlnet/args.py
@@ -471,3 +471,7 @@ def parse(cls, text: str) -> ControlNetUnit:
                 for (key, value) in (item.strip().split(": "),)
             },
         )
+
+    def __copy__(self) -> ControlNetUnit:
+        """Override the behavior on `copy.copy` calls."""
+        return self.copy()
diff --git a/unit_tests/args_test.py b/unit_tests/args_test.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 import numpy as np
+from copy import copy
 from dataclasses import dataclass
 
 from internal_controlnet.args import ControlNetUnit
@@ -249,3 +250,11 @@ def test_infotext_parsing():
 
 def test_alias():
     ControlNetUnit.from_dict({"lowvram": True})
+
+
+def test_copy():
+    unit1 = ControlNetUnit(enabled=True, module="none")
+    unit2 = copy(unit1)
+    unit2.enabled = False
+    assert unit1.enabled
+    assert not unit2.enabled