diff --git a/__init__.py b/__init__.py index 22c3b96..9a64b65 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,14 @@ settings.load_settings() if settings.SETTINGS.btp_enabled: from .py import betterTaesdPreview # noqa: F401 -from .py import deepshrink, hypertile, modelPatchConditional, samplers, sigmas +from .py import ( + deepshrink, + hypertile, + modelPatchConditional, + refinerAfter, + samplers, + sigmas, +) NODE_CLASS_MAPPINGS = { "BlehHyperTile": hypertile.HyperTileBleh, @@ -14,6 +21,7 @@ NODE_CLASS_MAPPINGS = { "BlehInsaneChainSampler": samplers.BlehInsaneChainSampler, "BlehForceSeedSampler": samplers.BlehForceSeedSampler, "BlehModelPatchConditional": modelPatchConditional.ModelPatchConditionalNode, + "BlehRefinerAfter": refinerAfter.BlehRefinerAfter, } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/py/refinerAfter.py b/py/refinerAfter.py new file mode 100644 index 0000000..e242b44 --- /dev/null +++ b/py/refinerAfter.py @@ -0,0 +1,49 @@ +from __future__ import annotations + + +class BlehRefinerAfter: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "start_timestep": ("INT", {"default": 199, "min": 0, "max": 1000}), + "model": ("MODEL",), + "refiner_model": ("MODEL",), + }, + } + + RETURN_TYPES = ("MODEL",) + CATEGORY = "bleh/model_patches" + + FUNCTION = "patch" + + @staticmethod + def get_real_model(model: object) -> object: + while hasattr(model, "model"): + model = model.model + return model + + def patch( + self, + start_timestep: int, + model: object, + refiner_model: object, + ) -> tuple[object]: + import comfy + + model = model.clone() + comfy.model_management.load_models_gpu([refiner_model]) + refiner_model = self.get_real_model(refiner_model) + ms = self.get_real_model(model).model_sampling + + def unet_wrapper(apply_model, args): + inp, timestep, c = args["input"], args["timestep"], args["c"] + curr_timestep = ms.timestep(timestep).float().max().item() + return ( + refiner_model.apply_model + if curr_timestep <= start_timestep + else apply_model + )(inp, timestep, **c) + + model.set_model_unet_function_wrapper(unet_wrapper) + return (model,)