Add BlehRefinerAfter node

__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 = {

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,)