46 lines
2.2 KiB
Diff
46 lines
2.2 KiB
Diff
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
|
|
index f116efee..e52466b3 100644
|
|
--- a/comfy/ldm/modules/attention.py
|
|
+++ b/comfy/ldm/modules/attention.py
|
|
@@ -460,6 +460,8 @@ class BasicTransformerBlock(nn.Module):
|
|
return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint)
|
|
|
|
def _forward(self, x, context=None, transformer_options={}):
|
|
+ if "transformer_function" in transformer_options:
|
|
+ return transformer_options["transformer_function"](self, x, context, transformer_options)
|
|
extra_options = {}
|
|
block = transformer_options.get("block", None)
|
|
block_index = transformer_options.get("block_index", 0)
|
|
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
|
|
index cf51c4ad..21d5b1af 100644
|
|
--- a/comfy/model_patcher.py
|
|
+++ b/comfy/model_patcher.py
|
|
@@ -158,6 +158,14 @@ class ModelPatcher:
|
|
def set_model_output_block_patch(self, patch):
|
|
self.set_model_patch(patch, "output_block_patch")
|
|
|
|
+ def set_model_transformer_function(self, transformer_function):
|
|
+ self.model_options["transformer_options"]["transformer_function"] = transformer_function
|
|
+
|
|
+ def set_model_sampler_cfg_rescaler(self, cfg_rescaler_fn, disable_cfg1_optimization=False):
|
|
+ self.model_options["sampler_cfg_rescaler"] = cfg_rescaler_fn
|
|
+ if disable_cfg1_optimization:
|
|
+ self.model_options["disable_cfg1_optimization"] = True
|
|
+
|
|
def add_object_patch(self, name, obj):
|
|
self.object_patches[name] = obj
|
|
|
|
diff --git a/comfy/samplers.py b/comfy/samplers.py
|
|
index 415a35cc..bd8912bc 100644
|
|
--- a/comfy/samplers.py
|
|
+++ b/comfy/samplers.py
|
|
@@ -249,6 +249,8 @@ def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_o
|
|
#The main sampling function shared by all the samplers
|
|
#Returns denoised
|
|
def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
|
|
+ if "sampler_cfg_rescaler" in model_options:
|
|
+ cond_scale = model_options["sampler_cfg_rescaler"]({"cond_scale": cond_scale, "timestep": timestep})
|
|
if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
|
|
uncond_ = None
|
|
else:
|