diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 84f582c..98877c9 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -69,7 +69,14 @@ class ModelPatcherHelper: self.remove_motion_models() self.remove_forward_timestep_embed_patch() - def get_motion_models(self) -> list[MotionModelPatcher]: + def get_motion_models(self, all_devices=False) -> list[MotionModelPatcher]: + if all_devices: + patchers = [self.model] + patchers.extend(self.model.get_additional_models_with_key("multigpu")) + all_motion_models = [] + for patcher in patchers: + all_motion_models.extend(patcher.get_additional_models_with_key(self.ADE)) + return all_motion_models return self.model.additional_models.get(self.ADE, []) def set_motion_models(self, motion_models: list[MotionModelPatcher]): @@ -82,7 +89,7 @@ class ModelPatcherHelper: self.model.remove_injections(self.ADE) def cleanup_motion_models(self): - for motion_model in self.get_motion_models(): + for motion_model in self.get_motion_models(all_devices=True): motion_model.cleanup() @@ -108,7 +115,7 @@ class ModelPatcherHelper: ########################## # motion models helpers def set_video_length(self, video_length: int, full_length: int): - for motion_model in self.get_motion_models(): + for motion_model in self.get_motion_models(all_devices=True): motion_model.model.set_video_length(video_length=video_length, full_length=full_length) def get_name_string(self, show_version=False): @@ -164,7 +171,7 @@ class ModelPatcherHelper: def pre_run(self): # TODO: could implement this as a ModelPatcher ON_PRE_RUN callback - for motion_model in self.get_motion_models(): + for motion_model in self.get_motion_models(all_devices=True): motion_model.pre_run() self.get_sample_settings().pre_run(self.model) @@ -670,7 +677,7 @@ class MotionModelAttachment: class MotionModelGroup: - def __init__(self, init_motion_model: MotionModelPatcher=None): + def __init__(self, init_motion_model: Union[MotionModelPatcher, list[MotionModelPatcher]]=None): self.models: list[MotionModelPatcher] = [] if init_motion_model is not None: if isinstance(init_motion_model, list): diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 2209a3c..3f57132 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -40,57 +40,82 @@ class AnimateDiffGlobalState: def __init__(self): self.model_patcher: ModelPatcher = None self.motion_models: MotionModelGroup = None + + self.model_patcher_devices: dict[torch.device, ModelPatcher] = {} + self.motion_models_devices: dict[torch.device, MotionModelGroup] = {} + self.params: InjectionParams = None self.sample_settings: SampleSettings = None self.callback_output_dict: dict[str] = {} self.function_injections: FunctionInjectionHolder = None self.reset() - def initialize(self, model: BaseModel): + def initialize(self, model: BaseModel, model_options: dict[str]): # this function is to be run in sampling func if not self.initialized: self.initialized = True - if self.motion_models is not None: - self.motion_models.initialize_timesteps(model) + # initialize multigpu stuff + if "multigpu_clones" in model_options: + self.model_patcher_devices = model_options["multigpu_clones"].copy() + for device, patcher in self.model_patcher_devices.items(): + mm_list = ModelPatcherHelper(patcher).get_motion_models() + if len(mm_list) > 0: + self.motion_models_devices[device] = MotionModelGroup(mm_list) + else: + self.model_patcher_devices[self.model_patcher.load_device] = self.model_patcher + if self.motion_models is not None: + self.motion_models_devices[self.model_patcher.load_device] = self.motion_models + # initialize timesteps + for device, motion_models in self.motion_models_devices.items(): + base_model = self.model_patcher_devices[device].model + motion_models.initialize_timesteps(base_model) if self.params.context_options is not None: self.params.context_options.initialize_timesteps(model) if self.sample_settings.custom_cfg is not None: self.sample_settings.custom_cfg.initialize_timesteps(model) def prepare_current_keyframes(self, x: Tensor, timestep: Tensor, transformer_options: dict[str, Tensor]): - if self.motion_models is not None: - self.motion_models.prepare_current_keyframe(x=x, t=timestep, transformer_options=transformer_options) + for motion_models in self.motion_models_devices.values(): + motion_models.prepare_current_keyframe(x=x, t=timestep, transformer_options=transformer_options) if self.params.context_options is not None: self.params.context_options.prepare_current(t=timestep, transformer_options=transformer_options) if self.sample_settings.custom_cfg is not None: self.sample_settings.custom_cfg.prepare_current_keyframe(t=timestep, transformer_options=transformer_options) - def perform_special_model_features(self, model: BaseModel, conds: list, x_in: Tensor, model_options: dict[str]): - if self.motion_models is not None: - special_models = self.motion_models.get_special_models() + def perform_special_model_features(self, conds: list, x_in: Tensor, model_options: dict[str]): + # despite there be + clone_uuids = set() + for device, motion_models in self.motion_models_devices.items(): + model: BaseModel = self.model_patcher_devices[device].model + special_models = motion_models.get_special_models() if len(special_models) > 0: for special_model in special_models: if special_model.model.is_in_effect(): attachment = get_mm_attachment(special_model) if attachment.is_pia(special_model): special_model.model.inject_unet_conv_in_pia_fancyvideo(model) - conds = get_conds_with_c_concat(conds, - attachment.get_pia_c_concat(model, x_in)) + if special_model.clone_base_uuid not in clone_uuids: + clone_uuids.add(special_model.clone_base_uuid) + conds = get_conds_with_c_concat(conds, + attachment.get_pia_c_concat(model, x_in)) elif attachment.is_fancyvideo(special_model): # TODO: handle other weights special_model.model.inject_unet_conv_in_pia_fancyvideo(model) - conds = get_conds_with_c_concat(conds, - attachment.get_fancy_c_concat(model, x_in)) - # add fps_embedding/motion_embedding patches - emb_patches = special_model.model.get_fancyvideo_emb_patches(dtype=x_in.dtype, device=x_in.device) - transformer_patches = model_options["transformer_options"].get("patches", {}) - transformer_patches["emb_patch"] = emb_patches - model_options["transformer_options"]["patches"] = transformer_patches + if special_model.clone_base_uuid not in clone_uuids: + clone_uuids.add(special_model.clone_base_uuid) + conds = get_conds_with_c_concat(conds, + attachment.get_fancy_c_concat(model, x_in)) + # add fps_embedding/motion_embedding patches + emb_patches = special_model.model.get_fancyvideo_emb_patches(dtype=x_in.dtype, device=x_in.device) + transformer_patches = model_options["transformer_options"].get("patches", {}) + transformer_patches["emb_patch"] = emb_patches + model_options["transformer_options"]["patches"] = transformer_patches return conds - def restore_special_model_features(self, model: BaseModel): - if self.motion_models is not None: - special_models = self.motion_models.get_special_models() + def restore_special_model_features(self): + for device, motion_models in self.motion_models_devices.items(): + model: BaseModel = self.model_patcher_devices[device].model + special_models = motion_models.get_special_models() if len(special_models) > 0: for special_model in reversed(special_models): attachment = get_mm_attachment(special_model) @@ -109,13 +134,17 @@ class AnimateDiffGlobalState: self.total_steps: int = 0 self.callback_output_dict.clear() self.callback_output_dict = {} + # model patchers + self.model_patcher_devices.clear() if self.model_patcher is not None: - self.model_patcher.clean_hooks() del self.model_patcher self.model_patcher = None + # motion models + self.motion_models_devices.clear() if self.motion_models is not None: del self.motion_models self.motion_models = None + # other if self.params is not None: self.params.context_options.reset() del self.params @@ -191,12 +220,13 @@ def _apply_model_wrapper(executor, *args, **kwargs): cond_or_uncond = transformer_options["cond_or_uncond"] ad_params = transformer_options["ad_params"] ADGS: AnimateDiffGlobalState = transformer_options["ADGS"] - if ADGS.motion_models is not None: - for motion_model in ADGS.motion_models.models: - attachment = get_mm_attachment(motion_model) - attachment.prepare_alcmi2v_features(motion_model, x=x, cond_or_uncond=cond_or_uncond, ad_params=ad_params, latent_format=executor.class_obj.latent_format) - attachment.prepare_camera_features(motion_model, x=x, cond_or_uncond=cond_or_uncond, ad_params=ad_params) - attachment.prepare_motionctrl_camera(motion_model, x=x, transformer_options=transformer_options) + motion_models = ADGS.motion_models_devices.get(x.device, None) + if motion_models is not None: + for motion_model in motion_models.models: + attachment = get_mm_attachment(motion_model) + attachment.prepare_alcmi2v_features(motion_model, x=x, cond_or_uncond=cond_or_uncond, ad_params=ad_params, latent_format=executor.class_obj.latent_format) + attachment.prepare_camera_features(motion_model, x=x, cond_or_uncond=cond_or_uncond, ad_params=ad_params) + attachment.prepare_motionctrl_camera(motion_model, x=x, transformer_options=transformer_options) del x return executor(*args, **kwargs) @@ -545,7 +575,7 @@ def outer_sample_wrapper(executor: WrapperExecutor, *args, **kwargs): def evolved_sampling_function(model, x: Tensor, timestep: Tensor, uncond, cond, cond_scale, model_options: dict={}, seed=None): ADGS: AnimateDiffGlobalState = model_options["transformer_options"]["ADGS"] - ADGS.initialize(model) + ADGS.initialize(model, model_options) ADGS.prepare_current_keyframes(x=x, timestep=timestep, transformer_options=model_options["transformer_options"]) try: # add AD/evolved-sampling params to model_options (transformer_options) @@ -556,7 +586,7 @@ def evolved_sampling_function(model, x: Tensor, timestep: Tensor, uncond, cond, model_options["transformer_options"] = model_options["transformer_options"].copy() model_options["transformer_options"]["ad_params"] = ADGS.create_exposed_params() - cond, uncond = ADGS.perform_special_model_features(model, [cond, uncond], x, model_options) + cond, uncond = ADGS.perform_special_model_features([cond, uncond], x, model_options) # only use cfg1_optimization if not using custom_cfg or explicitly set to 1.0 uncond_ = uncond @@ -576,7 +606,7 @@ def evolved_sampling_function(model, x: Tensor, timestep: Tensor, uncond, cond, return comfy.samplers.cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options, cond, uncond) finally: - ADGS.restore_special_model_features(model) + ADGS.restore_special_model_features() def perform_image_injection(ADGS: AnimateDiffGlobalState, model: BaseModel, latents: Tensor, to_inject: NoisedImageToInject) -> Tensor: @@ -634,9 +664,10 @@ def prepare_control_objects(control: ControlBase, full_idxs: list[int], ADGS: An if not hasattr(control, "sub_idxs"): raise ValueError(f"Control type {type(control).__name__} may not support required features for sliding context window; \ use ControlNet nodes from Kosinkadink/ComfyUI-Advanced-ControlNet, or make sure ComfyUI-Advanced-ControlNet is updated.") - control.sub_idxs = full_idxs - control.full_latent_length = ADGS.params.full_length - control.context_length = ADGS.params.context_options.context_length + if not hasattr(control, "ACN_VERSION"): + control.sub_idxs = full_idxs + control.full_latent_length = ADGS.params.full_length + control.context_length = ADGS.params.context_options.context_length # initial sliding_calc_conds_batch inspired by ashen's initial hack for 16-frame sliding context: # https://github.com/comfyanonymous/ComfyUI/compare/master...ashen-sensored:ComfyUI:master @@ -771,16 +802,16 @@ def sliding_calc_cond_batch(executor: Callable, model, conds: list[list[dict]], return conds_final ContextResults = collections.namedtuple("ContextResults", ['window_idx', 'sub_conds_out', 'sub_conds', 'ctx_idxs']) -def evaluate_context_windows(executor, model, x_in, conds, timestep, enumerated_context_windows: list[tuple[int, list[int]]], +def evaluate_context_windows(executor, model: BaseModel, x_in: Tensor, conds, timestep: Tensor, enumerated_context_windows: list[tuple[int, list[int]]], model_options, CREF: ContextRefHandler, ADGS: AnimateDiffGlobalState): results: list[ContextResults] = [] for window_idx, ctx_idxs in enumerated_context_windows: # allow processing to end between context window executions for faster Cancel comfy.model_management.throw_exception_if_processing_interrupted() ADGS.params.sub_idxs = ctx_idxs - if ADGS.motion_models is not None: - ADGS.motion_models.set_sub_idxs(ctx_idxs) - ADGS.motion_models.set_video_length(len(ctx_idxs), ADGS.params.full_length) + for motion_models in ADGS.motion_models_devices.values(): + motion_models.set_sub_idxs(ctx_idxs) + motion_models.set_video_length(len(ctx_idxs), ADGS.params.full_length) # update exposed params model_options["transformer_options"]["ad_params"]["sub_idxs"] = ctx_idxs model_options["transformer_options"]["ad_params"]["context_length"] = len(ctx_idxs) @@ -800,7 +831,6 @@ def combine_context_window_results(x_in: Tensor, sub_conds_out, sub_conds, ctx_i ADGS: AnimateDiffGlobalState, NAIVE: NaiveReuseHandler, CREF: ContextRefHandler, conds_final: list[Tensor], counts_final: list[Tensor], biases_final: list[Tensor]): if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE: - full_length = ADGS.params.full_length for pos, idx in enumerate(ctx_idxs): # bias is the influence of a specific index in relation to the whole context window bias = 1 - abs(idx - (ctx_idxs[0] + ctx_idxs[-1]) / 2) / ((ctx_idxs[-1] - ctx_idxs[0] + 1e-2) / 2)