mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-06-02 01:22:11 +08:00
65da29aaa9
* Make torch compile node use wrapper instead of object_patch for the entire diffusion_models object, allowing key assotiations on diffusion_models to not break (loras, getting attributes, etc.) * Moved torch compile code into comfy_api so it can be used by custom nodes with a degree of confidence * Refactor set_torch_compile_wrapper to support a list of keys instead of just diffusion_model, as well as additional torch.compile args * remove unused import * Moved torch compile kwargs to be stored in model_options instead of attachments; attachments are more intended for things to be 'persisted', AKA not deepcopied * Add some comments * Remove random line of code, not sure how it got there
24 lines
609 B
Python
24 lines
609 B
Python
from comfy_api.torch_helpers import set_torch_compile_wrapper
|
|
|
|
|
|
class TorchCompileModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model": ("MODEL",),
|
|
"backend": (["inductor", "cudagraphs"],),
|
|
}}
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
|
|
CATEGORY = "_for_testing"
|
|
EXPERIMENTAL = True
|
|
|
|
def patch(self, model, backend):
|
|
m = model.clone()
|
|
set_torch_compile_wrapper(model=m, backend=backend)
|
|
return (m, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TorchCompileModel": TorchCompileModel,
|
|
}
|