mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 14:50:49 +08:00
3c9d311dee
- the model weights generally have to be patched ahead of time for compilation to work - the model downloader matches the folder_paths API a bit better - tweak the logging from the execution node
79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Callable, Optional
|
|
|
|
import torch
|
|
|
|
import comfy.utils
|
|
from comfy import model_management
|
|
from comfy.patcher_extension import WrappersMP
|
|
|
|
if TYPE_CHECKING:
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.patcher_extension import WrapperExecutor
|
|
|
|
COMPILE_KEY = "torch.compile"
|
|
TORCH_COMPILE_KWARGS = "torch_compile_kwargs"
|
|
|
|
|
|
def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable], model_patcher: Optional[ModelPatcher] = None) -> Callable:
|
|
'''
|
|
Create a wrapper that will refer to the compiled_diffusion_model.
|
|
'''
|
|
|
|
def apply_torch_compile_wrapper(executor: WrapperExecutor, *args, **kwargs):
|
|
try:
|
|
orig_modules = {}
|
|
for key, value in compiled_module_dict.items():
|
|
orig_modules[key] = comfy.utils.get_attr(executor.class_obj, key)
|
|
comfy.utils.set_attr(executor.class_obj, key, value)
|
|
# todo: compilation has to patch all weights
|
|
if model_patcher is not None:
|
|
model_patcher.patch_model(device_to=model_management.get_torch_device(), force_patch_weights=True)
|
|
return executor(*args, **kwargs)
|
|
finally:
|
|
for key, value in orig_modules.items():
|
|
comfy.utils.set_attr(executor.class_obj, key, value)
|
|
|
|
return apply_torch_compile_wrapper
|
|
|
|
|
|
def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Optional[dict[str, str]] = None,
|
|
mode: Optional[str] = None, fullgraph=False, dynamic: Optional[bool] = None,
|
|
keys: list[str] = ["diffusion_model"], *args, **kwargs):
|
|
'''
|
|
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 modules.
|
|
'''
|
|
# clear out any other torch.compile wrappers
|
|
model.remove_wrappers_with_key(WrappersMP.APPLY_MODEL, COMPILE_KEY)
|
|
# if no keys, default to 'diffusion_model'
|
|
if not keys:
|
|
keys = ["diffusion_model"]
|
|
# create kwargs dict that can be referenced later
|
|
compile_kwargs = {
|
|
"backend": backend,
|
|
"options": options,
|
|
"mode": mode,
|
|
"fullgraph": fullgraph,
|
|
"dynamic": dynamic,
|
|
}
|
|
# get a dict of compiled keys
|
|
compiled_modules = {}
|
|
for key in keys:
|
|
compiled_modules[key] = torch.compile(
|
|
model=model.get_model_object(key),
|
|
**compile_kwargs,
|
|
)
|
|
# add torch.compile wrapper
|
|
wrapper_func = apply_torch_compile_factory(
|
|
compiled_module_dict=compiled_modules,
|
|
model_patcher=model,
|
|
)
|
|
# 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
|