From c64281ac95c7ad693dd953c68f57ef1fb9085f03 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 20 May 2025 22:32:03 -0700 Subject: [PATCH] Add some comments --- comfy_api/torch_helpers/torch_compile.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/comfy_api/torch_helpers/torch_compile.py b/comfy_api/torch_helpers/torch_compile.py index ebf977367..0dcc2e80c 100644 --- a/comfy_api/torch_helpers/torch_compile.py +++ b/comfy_api/torch_helpers/torch_compile.py @@ -15,7 +15,7 @@ TORCH_COMPILE_KWARGS = "torch_compile_kwargs" def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable]) -> Callable: ''' - Create a wrapper that will refer to the compiled_diffusion_model + Create a wrapper that will refer to the compiled_diffusion_model. ''' def apply_torch_compile_wrapper(executor: WrapperExecutor, *args, **kwargs): try: @@ -37,7 +37,7 @@ def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Option Perform torch.compile that will be applied at sample time for either the whole model or specific params of the BaseModel instance. When keys is None, it will default to using ["diffusion_model"], compiling the whole diffusion_model. - When a list of keys is provided, it will perform torch.compile on only the selected params. + When a list of keys is provided, it will perform torch.compile on only the selected modules. ''' torch.compile() # clear out any other torch.compile wrappers @@ -64,5 +64,7 @@ def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Option wrapper_func = apply_torch_compile_factory( compiled_module_dict=compiled_modules, ) + # store wrapper to run on BaseModel's apply_model function model.add_wrapper_with_key(WrappersMP.APPLY_MODEL, COMPILE_KEY, wrapper_func) + # keep compile kwargs for reference model.model_options[TORCH_COMPILE_KWARGS] = compile_kwargs