From bacb0b5eda05ca2fc3ca3e201034672c4a1490e0 Mon Sep 17 00:00:00 2001 From: Yo1up Date: Sun, 29 Jun 2025 20:45:34 -0400 Subject: [PATCH] Initial commit with basic model wrapper and unwrapper nodes --- Detail_Model_Patch_Custom.py | 169 +++++++++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100755 Detail_Model_Patch_Custom.py diff --git a/Detail_Model_Patch_Custom.py b/Detail_Model_Patch_Custom.py new file mode 100755 index 0000000..3fbcdbb --- /dev/null +++ b/Detail_Model_Patch_Custom.py @@ -0,0 +1,169 @@ +import torch +import torch.nn as nn + +class ModelWrapper(nn.Module): + """ + A wrapper class that intercepts model calls and allows custom processing + """ + def __init__(self, original_model, scaling_factor, narrowness): + super().__init__() + self.original_model = original_model + self.narrowness = narrowness + self.scaling_factor = scaling_factor + self._is_wrapped = True + + def unwrap(self): + return self.original_model + + def sin_schedule(self, x): + new_x = x / 1000 # (x - self.sigma_min) / (self.sigma_max - self.sigma_min) + return torch.sin(new_x * torch.pi) ** self.narrowness + + def forward(self, *args, **kwargs): + # PRE-PROCESSING SPACE - Add your custom code here + # You can modify inputs, log information, etc. + # Example: print(f"Model called with args: {len(args)}, kwargs: {list(kwargs.keys())}") + + # Call the original model + result = self.original_model(*args, **kwargs) + + print(args[1]) + timestep = args[1][0] + print(timestep) + result = result + (result * self.sin_schedule(timestep) * self.scaling_factor) + # POST-PROCESSING SPACE - Add your custom code here + # You can modify outputs, save results, apply transformations, etc. + # Example: print(f"Model output shape: {result.shape if hasattr(result, 'shape') else type(result)}") + + return result + + def __getattribute__(self, name): + # Always use direct access for internal stuff + if name in {"original_model", "scaling_factor", "narrowness", "_is_wrapped", "unwrap"}: + return object.__getattribute__(self, name) + try: + return object.__getattribute__(self, name) + except AttributeError: + # fallback to wrapped model + return getattr(self.original_model, name) + + + +class DetailModelWrapperNode: + """ + A ComfyUI node that wraps a model with custom processing capabilities + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL", ), + "scaling_factor": ("FLOAT", {"default": 0.05, "min": -1.0, "max": 1.0, "step": 0.005}), + "scaling_narrowness": ("FLOAT", ), + # "effect_depth": ("FLOAT", {"default": 0.5, "min": 0, "max": 1.0, "step": 0.05}) + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("wrapped_model",) + FUNCTION = "wrap_model" + CATEGORY = "model_processing" + + def unwrap_model(self, wrapped_model): + """ + Unwraps the input model with custom processing capabilities + """ + + unwrapped_model = wrapped_model.clone() + + # Get the actual model object (usually under 'model' attribute) + if hasattr(wrapped_model, 'model') and hasattr(wrapped_model.model, 'diffusion_model') and hasattr(wrapped_model.model.diffusion_model, "unwrap"): + # Wrap the diffusion model (for SD models) + unwrapped_model.model.diffusion_model = wrapped_model.model.diffusion_model.unwrap() + elif hasattr(wrapped_model, 'model') and hasattr(wrapped_model.model, "unwrap"): + # Wrap the main model + unwrapped_model.model = wrapped_model.model.unwrap() + elif hasattr(wrapped_model, "unwrap"): + # Unwrap the whole object + unwrapped_model = wrapped_model.unwrap() + else: + # Fallback to passing model straight through if no unwrapper can be found. + unwrapped_model = wrapped_model + + return (unwrapped_model,) + + def wrap_model(self, model, scaling_factor, scaling_narrowness): + """ + Wraps the input model with custom processing capabilities + """ + + model = self.unwrap_model(model.clone()) + + # Clone the model to avoid modifying the original + wrapped_model = model[0].clone() + + # Get the actual model object (usually under 'model' attribute) + if hasattr(wrapped_model, 'model') and hasattr(wrapped_model.model, 'diffusion_model'): + # Wrap the diffusion model (for SD models) + wrapped_model.model.diffusion_model = ModelWrapper(wrapped_model.model.diffusion_model, scaling_factor, scaling_narrowness) + elif hasattr(wrapped_model, 'model'): + # Wrap the main model + wrapped_model.model = ModelWrapper(wrapped_model.model, scaling_factor, scaling_narrowness) + else: + # Fallback: wrap the entire model object + wrapped_model = ModelWrapper(wrapped_model, scaling_factor, scaling_narrowness) + + return (wrapped_model,) + +class DetailModelUnwrapperNode: + """ + A ComfyUI node that wraps a model with custom processing capabilities + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL", ), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("unwrapped_model",) + FUNCTION = "unwrap_model" + CATEGORY = "model_processing" + + def unwrap_model(self, model): + """ + Unwraps the input model with custom processing capabilities + """ + + unwrapped_model = model.clone() + + # Get the actual model object (usually under 'model' attribute) + if hasattr(model, 'model') and hasattr(model.model, 'diffusion_model') and hasattr(model.model.diffusion_model, "unwrap"): + # Wrap the diffusion model (for SD models) + unwrapped_model.model.diffusion_model = model.model.diffusion_model.unwrap() + elif hasattr(model, 'model') and hasattr(model.model, "unwrap"): + # Wrap the main model + unwrapped_model.model = model.model.unwrap() + elif hasattr(model, "unwrap"): + # Unwrap the whole object + unwrapped_model = model.unwrap() + else: + # Fallback to passing model straight through if no unwrapper can be found. + unwrapped_model = model + + return (unwrapped_model,) + +# Node registration for ComfyUI +NODE_CLASS_MAPPINGS = { + "DetailModelWrapperNode": DetailModelWrapperNode, + "DetailModelUnwrapperNode": DetailModelUnwrapperNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DetailModelWrapperNode": "Detail Model Wrapper", + "DetailModelUnwrapperNode": "Detail Model Unwrapper" +}