Merge PR #169 from Kosinkadink/develop - initial flux support

Initial flux support, refactoring weight control

Author: Kosinkadink

adv_control/control.py

@@ -8,7 +8,7 @@
 import comfy.model_management
 import comfy.model_detection
 import comfy.controlnet as comfy_cn
-from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter
+from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter, StrengthType
 from comfy.model_patcher import ModelPatcher
 
 from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseSettings, SparseConst
@@ -21,13 +21,23 @@
 
 
 class ControlNetAdvanced(ControlNet, AdvancedControlBase):
-    def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None):
+    def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT):
         super().__init__(control_model=control_model, global_average_pooling=global_average_pooling, compression_ratio=compression_ratio, latent_format=latent_format, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
         AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet())
+        self.is_flux = False
+        self.x_noisy_shape = None
 
     def get_universal_weights(self) -> ControlWeights:
-        raw_weights = [(self.weights.base_multiplier ** float(12 - i)) for i in range(13)]
-        return self.weights.copy_with_new_weights(raw_weights)
+        def cn_weights_func(idx: int, control: dict[str, list[Tensor]], key: str):
+            if key == "middle":
+                return 1.0
+            c_len = len(control[key])
+            raw_weights = [(self.weights.base_multiplier ** float((c_len) - i)) for i in range(c_len+1)]
+            raw_weights = raw_weights[:-1]
+            if key == "input":
+                raw_weights.reverse()
+            return raw_weights[idx]
+        return self.weights.copy_with_new_weights(new_weight_func=cn_weights_func)
 
     def get_control_advanced(self, x_noisy, t, cond, batched_number):
         # perform special version of get_control that supports sliding context and masks
@@ -49,7 +59,6 @@ def sliding_get_control(self, x_noisy: Tensor, t, cond, batched_number):
         if self.manual_cast_dtype is not None:
             dtype = self.manual_cast_dtype
 
-        output_dtype = x_noisy.dtype
         # make cond_hint appropriate dimensions
         # TODO: change this to not require cond_hint upscaling every step when self.sub_idxs are present
         if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:
@@ -64,9 +73,9 @@ def sliding_get_control(self, x_noisy: Tensor, t, cond, batched_number):
                 actual_cond_hint_orig = self.cond_hint_original
                 if self.cond_hint_original.size(0) < self.full_latent_length:
                     actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length)
-                self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
+                self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
             else:
-                self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
+                self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
             if self.vae is not None:
                 loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
                 self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1))
@@ -81,25 +90,44 @@ def sliding_get_control(self, x_noisy: Tensor, t, cond, batched_number):
         self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
 
         context = cond.get('crossattn_controlnet', cond['c_crossattn'])
-        y = cond.get('y', None)
-        if y is not None:
-            y = y.to(dtype)
+        extra = self.extra_args.copy()
+        for c in self.extra_conds:
+            temp = cond.get(c, None)
+            if temp is not None:
+                extra[c] = temp.to(dtype)
+
         timestep = self.model_sampling_current.timestep(t)
         x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
+        self.x_noisy_shape = x_noisy.shape
+        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
+        return self.control_merge(control, control_prev, output_dtype=None)
 
-        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
-        return self.control_merge(control, control_prev, output_dtype)
+    def pre_run_advanced(self, *args, **kwargs):
+        self.is_flux = "Flux" in str(type(self.control_model).__name__)
+        return super().pre_run_advanced(*args, **kwargs)
+
+    def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int, flux_shape=None):
+        if self.is_flux:
+            flux_shape = self.x_noisy_shape
+        return super().apply_advanced_strengths_and_masks(x, batched_number, flux_shape)
 
     def copy(self):
         c = ControlNetAdvanced(self.control_model, self.timestep_keyframes, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
+        c.control_model = self.control_model
+        c.control_model_wrapped = self.control_model_wrapped
         self.copy_to(c)
         self.copy_to_advanced(c)
         return c
 
+    def cleanup_advanced(self):
+        self.x_noisy_shape = None
+        return super().cleanup_advanced()
+
     @staticmethod
     def from_vanilla(v: ControlNet, timestep_keyframe: TimestepKeyframeGroup=None) -> 'ControlNetAdvanced':
         to_return = ControlNetAdvanced(control_model=v.control_model, timestep_keyframes=timestep_keyframe,
-                                  global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device, manual_cast_dtype=v.manual_cast_dtype)
+                                  global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device,
+                                  manual_cast_dtype=v.manual_cast_dtype)
         v.copy_to(to_return)
         return to_return
 
@@ -121,18 +149,28 @@ def control_merge_inject(self, control: dict[str, list[Tensor]], control_prev, o
         return AdvancedControlBase.control_merge_inject(self, control, control_prev, output_dtype)
 
     def get_universal_weights(self) -> ControlWeights:
-        raw_weights = [(self.weights.base_multiplier ** float(7 - i)) for i in range(8)]
-        raw_weights = [raw_weights[-8], raw_weights[-3], raw_weights[-2], raw_weights[-1]]
-        raw_weights = get_properly_arranged_t2i_weights(raw_weights)
-        raw_weights.reverse()  # need to reverse to match recent ComfyUI changes
-        return self.weights.copy_with_new_weights(raw_weights)
+        def t2i_weights_func(idx: int, control: dict[str, list[Tensor]], key: str):
+            if key == "middle":
+                return 1.0
+            c_len = 8 #len(control[key])
+            raw_weights = [(self.weights.base_multiplier ** float((c_len-1) - i)) for i in range(c_len)]
+            raw_weights = [raw_weights[-c_len], raw_weights[-3], raw_weights[-2], raw_weights[-1]]
+            raw_weights = get_properly_arranged_t2i_weights(raw_weights)
+            if key == "input":
+                raw_weights.reverse()
+            return raw_weights[idx]
+        return self.weights.copy_with_new_weights(new_weight_func=t2i_weights_func)
 
     def get_calc_pow(self, idx: int, control: dict[str, list[Tensor]], key: str) -> int:
+        if key == "middle":
+            return 0
         # match how T2IAdapterAdvanced deals with universal weights
-        indeces = [7 - i for i in range(8)]
-        indeces = [indeces[-8], indeces[-3], indeces[-2], indeces[-1]]
+        c_len = 8 #len(control[key])
+        indeces = [(c_len-1) - i for i in range(c_len)]
+        indeces = [indeces[-c_len], indeces[-3], indeces[-2], indeces[-1]]
         indeces = get_properly_arranged_t2i_weights(indeces)
-        indeces.reverse()  # need to reverse to match recent ComfyUI changes
+        if key == "input":
+            indeces.reverse()  # need to reverse to match recent ComfyUI changes
         return indeces[idx]
 
     def get_control_advanced(self, x_noisy, t, cond, batched_number):
@@ -381,11 +419,11 @@ def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
         control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
         return self.control_merge(control, control_prev, output_dtype)
 
-    def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int):
+    def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int, *args, **kwargs):
         # apply mults to indexes with and without a direct condhint
         x[self.local_sparse_idxs] *= self.sparse_settings.sparse_hint_mult * self.weights.extras.get(SparseConst.HINT_MULT, 1.0)
         x[self.local_sparse_idxs_inverse] *= self.sparse_settings.sparse_nonhint_mult * self.weights.extras.get(SparseConst.NONHINT_MULT, 1.0)
-        return super().apply_advanced_strengths_and_masks(x, batched_number)
+        return super().apply_advanced_strengths_and_masks(x, batched_number, *args, **kwargs)
 
     def pre_run_advanced(self, model, percent_to_timestep_function):
         super().pre_run_advanced(model, percent_to_timestep_function)

adv_control/control_plusplus.py

@@ -237,9 +237,16 @@ def __init__(self, control_model: ControlNetPlusPlus, timestep_keyframes: Timest
         self.single_control_type: str = None
 
     def get_universal_weights(self) -> ControlWeights:
-        # TODO: match actual layer count of model
-        raw_weights = [(self.weights.base_multiplier ** float(12 - i)) for i in range(13)]
-        return self.weights.copy_with_new_weights(raw_weights)
+        def cn_weights_func(idx: int, control: dict[str, list[Tensor]], key: str):
+            if key == "middle":
+                return 1.0
+            c_len = len(control[key])
+            raw_weights = [(self.weights.base_multiplier ** float((c_len) - i)) for i in range(c_len+1)]
+            raw_weights = raw_weights[:-1]
+            if key == "input":
+                raw_weights.reverse()
+            return raw_weights[idx]
+        return self.weights.copy_with_new_weights(new_weight_func=cn_weights_func)
 
     def verify_control_type(self, model_name: str, pp_group: PlusPlusInputGroup=None):
         if pp_group is not None: