Merge PR #460 from Kosinkadink/develop - New ComfyUI version changes

Work with new ComfyUI ModelPatcher update (NOT backwards compatible)

animatediff/model_injection.py

@@ -220,53 +220,54 @@ def get_combined_hooked_patches(self, lora_hooks: LoraHookGroup):
                     combined_patches[key] = current_patches
         return combined_patches
 
-    def model_patches_to(self, device):
-        super().model_patches_to(device)
-
-    def patch_model(self, device_to=None, patch_weights=True):
+    def patch_model(self, *args, **kwargs):
+        was_injected = False
+        if self.currently_injected:
+            self.eject_model()
+            was_injected = True
         # first, perform model patching
-        if patch_weights: # TODO: keep only 'else' portion when don't need to worry about past comfy versions
-            patched_model = super().patch_model(device_to)
-        else:
-            patched_model = super().patch_model(device_to, patch_weights)
-        # finally, perform motion model injection
-        self.inject_model()
+        patched_model = super().patch_model(*args, **kwargs)
+        # bring injection back to original state
+        if was_injected and not self.currently_injected:
+            self.inject_model()
         return patched_model
 
-    def patch_model_lowvram(self, *args, **kwargs):
+    def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
+        self.eject_model()
         try:
-            return super().patch_model_lowvram(*args, **kwargs)
+            return super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
         finally:
-            # check if any modules have weight_function or bias_function that is not None
-            # NOTE: this serves no purpose currently, but I have it here for future reasons
-            for n, m in self.model.named_modules():
-                if not hasattr(m, "comfy_cast_weights"):
-                    continue
-                if getattr(m, "weight_function", None) is not None:
-                    self.model_params_lowvram = True
-                    self.model_params_lowvram_keys[f"{n}.weight"] = n
-                if getattr(m, "bias_function", None) is not None:
-                    self.model_params_lowvram = True
-                    self.model_params_lowvram_keys[f"{n}.bias"] = n
+            self.inject_model()
+            if lowvram_model_memory > 0:
+                self._patch_lowvram_extras()
+
+    def _patch_lowvram_extras(self):
+        # check if any modules have weight_function or bias_function that is not None
+        # NOTE: this serves no purpose currently, but I have it here for future reasons
+        self.model_params_lowvram = False
+        self.model_params_lowvram_keys.clear()
+        for n, m in self.model.named_modules():
+            if not hasattr(m, "comfy_cast_weights"):
+                continue
+            if getattr(m, "weight_function", None) is not None:
+                self.model_params_lowvram = True
+                self.model_params_lowvram_keys[f"{n}.weight"] = n
+            if getattr(m, "bias_function", None) is not None:
+                self.model_params_lowvram = True
+                self.model_params_lowvram_keys[f"{n}.bias"] = n  
 
     def unpatch_model(self, device_to=None, unpatch_weights=True):
         # first, eject motion model from unet
         self.eject_model()
         # finally, do normal model unpatching
-        if unpatch_weights: # TODO: keep only 'else' portion when don't need to worry about past comfy versions
+        if unpatch_weights:
             # handle hooked_patches first
             self.clean_hooks()
-            try:
-                return super().unpatch_model(device_to)
-            finally:
-                self.model_params_lowvram = False
-                self.model_params_lowvram_keys.clear()
-        else:
-            try:
-                return super().unpatch_model(device_to, unpatch_weights)
-            finally:
-                self.model_params_lowvram = False
-                self.model_params_lowvram_keys.clear()
+        try:
+            return super().unpatch_model(device_to, unpatch_weights)
+        finally:
+            self.model_params_lowvram = False
+            self.model_params_lowvram_keys.clear()
 
     def partially_load(self, *args, **kwargs):
         # partially_load calls patch_model, but we don't want to inject model in the intermediate call;
@@ -625,31 +626,37 @@ def patch_hooked_replace_weight_to_device(self, model_sd: dict, replace_patches:
             else:
                 comfy.utils.set_attr_param(self.model, key, out_weight)
 
-    def patch_model(self, device_to=None, patch_weights=True, *args, **kwargs):
+    def patch_model(self, device_to=None, *args, **kwargs):
         if self.desired_lora_hooks is not None:
             self.patches_backup = self.patches.copy()
             relevant_patches = self.get_combined_hooked_patches(lora_hooks=self.desired_lora_hooks)
             for key in relevant_patches:
                 self.patches.setdefault(key, [])
                 self.patches[key].extend(relevant_patches[key])
             self.current_lora_hooks = self.desired_lora_hooks
-        return super().patch_model(device_to, patch_weights, *args, **kwargs)
+        return super().patch_model(device_to, *args, **kwargs)
 
-    def patch_model_lowvram(self, *args, **kwargs):
+    def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
         try:
-            return super().patch_model_lowvram(*args, **kwargs)
+            return super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
         finally:
-            # check if any modules have weight_function or bias_function that is not None
-            # NOTE: this serves no purpose currently, but I have it here for future reasons
-            for n, m in self.model.named_modules():
-                if not hasattr(m, "comfy_cast_weights"):
-                    continue
-                if getattr(m, "weight_function", None) is not None:
-                    self.model_params_lowvram = True
-                    self.model_params_lowvram_keys[f"{n}.weight"] = n
-                if getattr(m, "bias_function", None) is not None:
-                    self.model_params_lowvram = True
-                    self.model_params_lowvram_keys[f"{n}.weight"] = n
+            if lowvram_model_memory > 0:
+                self._patch_lowvram_extras()
+
+    def _patch_lowvram_extras(self):
+        # check if any modules have weight_function or bias_function that is not None
+        # NOTE: this serves no purpose currently, but I have it here for future reasons
+        self.model_params_lowvram = False
+        self.model_params_lowvram_keys.clear()
+        for n, m in self.model.named_modules():
+            if not hasattr(m, "comfy_cast_weights"):
+                continue
+            if getattr(m, "weight_function", None) is not None:
+                self.model_params_lowvram = True
+                self.model_params_lowvram_keys[f"{n}.weight"] = n
+            if getattr(m, "bias_function", None) is not None:
+                self.model_params_lowvram = True
+                self.model_params_lowvram_keys[f"{n}.weight"] = n
 
     def unpatch_model(self, device_to=None, unpatch_weights=True, *args, **kwargs):
         try:
@@ -797,10 +804,14 @@ def __init__(self, *args, **kwargs):
         self.was_within_range = False
         self.prev_sub_idxs = None
         self.prev_batched_number = None
-
-    def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, *args, **kwargs):
-        patched_model = super().patch_model_lowvram(device_to, lowvram_model_memory, force_patch_weights, *args, **kwargs)
 
+    def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
+        to_return = super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
+        if lowvram_model_memory > 0:
+            self._patch_lowvram_extras(device_to=device_to)
+        return to_return
+
+    def _patch_lowvram_extras(self, device_to=None):
         # figure out the tensors (likely pe's) that should be cast to device besides just the named_modules
         remaining_tensors = list(self.model.state_dict().keys())
         named_modules = []
@@ -817,8 +828,6 @@ def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patc
             if device_to is not None:
                 comfy.utils.set_attr(self.model, key, comfy.utils.get_attr(self.model, key).to(device_to))
 
-        return patched_model
-
     def pre_run(self, model: ModelPatcherAndInjector):
         self.cleanup()
         self.model.set_scale(self.scale_multival, self.per_block_list)

animatediff/sampling.py

@@ -12,12 +12,7 @@
 from comfy.ldm.modules.diffusionmodules import openaimodel
 import comfy.model_management
 import comfy.samplers
-import comfy.sample
-SAMPLE_FALLBACK = False
-try:
-    import comfy.sampler_helpers
-except ImportError:
-    SAMPLE_FALLBACK = True
+import comfy.sampler_helpers
 import comfy.utils
 from comfy.controlnet import ControlBase
 from comfy.model_base import BaseModel
@@ -291,10 +286,7 @@ def inject_functions(self, model: ModelPatcherAndInjector, params: InjectionPara
         self.orig_diffusion_model_forward = model.model.diffusion_model.forward
         self.orig_sampling_function = comfy.samplers.sampling_function # used to support sliding context windows in samplers
         self.orig_get_area_and_mult = comfy.samplers.get_area_and_mult
-        if SAMPLE_FALLBACK:  # for backwards compatibility, for now
-            self.orig_get_additional_models = comfy.sample.get_additional_models
-        else:
-            self.orig_get_additional_models = comfy.sampler_helpers.get_additional_models
+        self.orig_get_additional_models = comfy.sampler_helpers.get_additional_models
         self.orig_apply_model = model.model.apply_model
         # Inject Functions
         openaimodel.forward_timestep_embed = forward_timestep_embed_factory()
@@ -324,10 +316,7 @@ def inject_functions(self, model: ModelPatcherAndInjector, params: InjectionPara
             del info
         comfy.samplers.sampling_function = evolved_sampling_function
         comfy.samplers.get_area_and_mult = get_area_and_mult_ADE
-        if SAMPLE_FALLBACK:  # for backwards compatibility, for now
-            comfy.sample.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_models)
-        else:
-            comfy.sampler_helpers.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_models)
+        comfy.sampler_helpers.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_models)
         # create temp_uninjector to help facilitate uninjecting functions
         self.temp_uninjector = GroupnormUninjectHelper(self)
 
@@ -341,10 +330,7 @@ def restore_functions(self, model: ModelPatcherAndInjector):
             model.model.diffusion_model.forward = self.orig_diffusion_model_forward
             comfy.samplers.sampling_function = self.orig_sampling_function
             comfy.samplers.get_area_and_mult = self.orig_get_area_and_mult
-            if SAMPLE_FALLBACK:  # for backwards compatibility, for now
-                comfy.sample.get_additional_models = self.orig_get_additional_models
-            else:
-                comfy.sampler_helpers.get_additional_models = self.orig_get_additional_models
+            comfy.sampler_helpers.get_additional_models = self.orig_get_additional_models
             model.model.apply_model = self.orig_apply_model
         except AttributeError:
             logger.error("Encountered AttributeError while attempting to restore functions - likely, an error occured while trying " + \
@@ -505,17 +491,8 @@ def ad_callback(step, x0, x, total_steps):
                 if is_custom:
                     iter_kwargs[IterationOptions.SAMPLER] = None #args[-5]
                 else:
-                    if SAMPLE_FALLBACK:  # backwards compatibility, for now
-                        # in older comfy, model needs to be loaded to get proper model_sampling to be used for sigmas
-                        comfy.model_management.load_model_gpu(model)
-                        iter_model = model.model
-                    else:
-                        iter_model = model
-                    current_device = None
-                    if hasattr(model, "current_device"): # backwards compatibility, for now
-                        current_device = model.current_device
-                    else:
-                        current_device = model.model.device
+                    iter_model = model
+                    current_device = model.model.device
                     iter_kwargs[IterationOptions.SAMPLER] = comfy.samplers.KSampler(
                         iter_model, steps=999, #steps=args[-7],
                         device=current_device, sampler=args[-5],
@@ -653,35 +630,20 @@ def evolved_sampling_function(model, x: Tensor, timestep: Tensor, uncond, cond,
         model_options["transformer_options"]["ad_params"] = ADGS.create_exposed_params()
 
         if not ADGS.is_using_sliding_context():
-            cond_pred, uncond_pred = calc_cond_uncond_batch_wrapper(model, [cond, uncond_], x, timestep, model_options)
+            cond_pred, uncond_pred = calc_conds_batch_wrapper(model, [cond, uncond_], x, timestep, model_options)
         else:
             cond_pred, uncond_pred = sliding_calc_conds_batch(model, [cond, uncond_], x, timestep, model_options)
 
-        if hasattr(comfy.samplers, "cfg_function"):
-            if ADGS.sample_settings.custom_cfg is not None:
-                cond_scale = ADGS.sample_settings.custom_cfg.get_cfg_scale(cond_pred)
-                model_options = ADGS.sample_settings.custom_cfg.get_model_options(model_options)
-            try:
-                cached_calc_cond_batch = comfy.samplers.calc_cond_batch
-                # support hooks and sliding context for PAG/other sampler_post_cfg_function tech that may use calc_cond_batch
-                comfy.samplers.calc_cond_batch = wrapped_cfg_sliding_calc_cond_batch_factory(cached_calc_cond_batch)
-                return comfy.samplers.cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options, cond, uncond)
-            finally:
-                comfy.samplers.calc_cond_batch = cached_calc_cond_batch
-        else: # for backwards compatibility, for now
-            if "sampler_cfg_function" in model_options:
-                args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
-                        "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
-                cfg_result = x - model_options["sampler_cfg_function"](args)
-            else:
-                cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
-
-            for fn in model_options.get("sampler_post_cfg_function", []):
-                args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
-                        "sigma": timestep, "model_options": model_options, "input": x}
-                cfg_result = fn(args)
-
-            return cfg_result
+        if ADGS.sample_settings.custom_cfg is not None:
+            cond_scale = ADGS.sample_settings.custom_cfg.get_cfg_scale(cond_pred)
+            model_options = ADGS.sample_settings.custom_cfg.get_model_options(model_options)
+        try:
+            cached_calc_cond_batch = comfy.samplers.calc_cond_batch
+            # support hooks and sliding context for PAG/other sampler_post_cfg_function tech that may use calc_cond_batch
+            comfy.samplers.calc_cond_batch = wrapped_cfg_sliding_calc_cond_batch_factory(cached_calc_cond_batch)
+            return comfy.samplers.cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options, cond, uncond)
+        finally:
+            comfy.samplers.calc_cond_batch = cached_calc_cond_batch
     finally:
         ADGS.restore_special_model_features(model)
 
@@ -745,7 +707,7 @@ def wrapped_cfg_sliding_calc_cond_batch(model, conds, x_in, timestep, model_opti
             # when inside sliding_calc_conds_batch, should return to original calc_cond_batch
             comfy.samplers.calc_cond_batch = orig_calc_cond_batch
             if not ADGS.is_using_sliding_context():
-                return calc_cond_uncond_batch_wrapper(model, conds, x_in, timestep, model_options)
+                return calc_conds_batch_wrapper(model, conds, x_in, timestep, model_options)
             else:
                 return sliding_calc_conds_batch(model, conds, x_in, timestep, model_options)
         finally:
@@ -922,7 +884,7 @@ def get_resized_cond(cond_in, full_idxs: list[int], context_length: int) -> list
                 model_options["transformer_options"][CONTEXTREF_MACHINE_STATE] = MachineState.OFF
             #logger.info(f"window: {curr_window_idx} - {model_options['transformer_options'][CONTEXTREF_MACHINE_STATE]}")
 
-            sub_conds_out = calc_cond_uncond_batch_wrapper(model, sub_conds, sub_x, sub_timestep, model_options)
+            sub_conds_out = calc_conds_batch_wrapper(model, sub_conds, sub_x, sub_timestep, model_options)
 
             if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE:
                 full_length = ADGS.params.full_length
@@ -1008,7 +970,7 @@ def get_conds_with_c_concat(conds: list[dict], c_concat: comfy.conds.CONDNoiseSh
     return new_conds
 
 
-def calc_cond_uncond_batch_wrapper(model, conds: list[dict], x_in: Tensor, timestep, model_options):
+def calc_conds_batch_wrapper(model, conds: list[dict], x_in: Tensor, timestep, model_options):
     # check if conds or unconds contain lora_hook or default_cond
     contains_lora_hooks = False
     has_default_cond = False
@@ -1028,9 +990,6 @@ def calc_cond_uncond_batch_wrapper(model, conds: list[dict], x_in: Tensor, times
        ADGS.hooks_initialize(model, hook_groups=hook_groups)
        ADGS.prepare_hooks_current_keyframes(timestep, hook_groups=hook_groups)
        return calc_conds_batch_lora_hook(model, conds, x_in, timestep, model_options, has_default_cond)
-    # keep for backwards compatibility, for now
-    if not hasattr(comfy.samplers, "calc_cond_batch"):
-        return comfy.samplers.calc_cond_uncond_batch(model, conds[0], conds[1], x_in, timestep, model_options)
     return comfy.samplers.calc_cond_batch(model, conds, x_in, timestep, model_options) 
 
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-animatediff-evolved"
 description = "Improved AnimateDiff integration for ComfyUI."
-version = "1.1.4"
+version = "1.2.0"
 license = { file = "LICENSE" }
 dependencies = []