From d1d7f322a8df42405b96837db948299cf0e5f97d Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 6 Feb 2024 04:30:54 -0600 Subject: [PATCH] Fixed ControlLLLite issue with tiled samplers --- adv_control/control.py | 4 ++-- adv_control/control_lllite.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index a140e5c..cd6783b 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -287,8 +287,8 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device=None): super().__init__(device) AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), require_model=True) - self.patch_attn1 = patch_attn1.clone_with_control(self) - self.patch_attn2 = patch_attn2.clone_with_control(self) + self.patch_attn1 = patch_attn1.set_control(self) + self.patch_attn2 = patch_attn2.set_control(self) self.latent_dims_div2 = None self.latent_dims_div4 = None diff --git a/adv_control/control_lllite.py b/adv_control/control_lllite.py index 0cdd01c..d451f91 100644 --- a/adv_control/control_lllite.py +++ b/adv_control/control_lllite.py @@ -82,8 +82,9 @@ class LLLitePatch: self.modules[d] = self.modules[d].to(device) return self - def set_control(self, control: Union[AdvancedControlBase, ControlBase]): + def set_control(self, control: Union[AdvancedControlBase, ControlBase]) -> 'LLLitePatch': self.control = control + return self #logger.error(f"set control for LLLitePatch: {id(self)}, cn: {id(control)}") def clone_with_control(self, control: AdvancedControlBase):