From 67be7eb81d59a7997e79a58e12e408e778d976ff Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 22 Jul 2023 17:01:12 -0400 Subject: [PATCH] Nodes can now patch the unet function. --- comfy/samplers.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 50fda016..9eee25a9 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -248,7 +248,10 @@ def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, con c['transformer_options'] = transformer_options - output = model_function(input_x, timestep_, **c).chunk(batch_chunks) + if 'model_function_wrapper' in model_options: + output = model_options['model_function_wrapper'](model_function, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks) + else: + output = model_function(input_x, timestep_, **c).chunk(batch_chunks) del input_x model_management.throw_exception_if_processing_interrupted()