19 lines
474 B
Python
19 lines
474 B
Python
|
|
from ..utils.module_utils import setup_ref_attn
|
|
|
|
class ApplyRefUNetNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"attn": (["FULL"],),
|
|
}}
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply"
|
|
|
|
CATEGORY = "flow"
|
|
|
|
def apply(self, model, attn):
|
|
ref_controller = setup_ref_attn(model, attn)
|
|
model.model_options['ref_controller'] = ref_controller
|
|
return (model, ) |