25 lines
861 B
Python
25 lines
861 B
Python
from typing import Callable
|
|
|
|
from comfy.model_patcher import ModelPatcher
|
|
from .util import is_injected_model, get_injected_model
|
|
|
|
|
|
def keyframe_sample_factory(orig_comfy_sample: Callable) -> Callable:
|
|
def keyframe_sample(model: ModelPatcher, *args, **kwargs):
|
|
if not is_injected_model(model.model):
|
|
return orig_comfy_sample(model, *args, **kwargs)
|
|
inject_param = get_injected_model(model.model)
|
|
try:
|
|
inject_param.reset()
|
|
|
|
inject_param.noise = args[0]
|
|
inject_param.steps = args[1]
|
|
inject_param.scheduler = args[4]
|
|
inject_param.denoise = kwargs.get('denoise', None)
|
|
inject_param.seed = kwargs.get('seed', None)
|
|
return orig_comfy_sample(model, *args, **kwargs)
|
|
finally:
|
|
inject_param.reset()
|
|
|
|
return keyframe_sample
|