1
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00

Add type annotation UnetWrapperFunction (#3531)

* Add type annotation UnetWrapperFunction

* nit

* Add types.py
This commit is contained in:
Chenlei Hu
2024-05-22 02:07:27 -04:00
committed by GitHub
parent 8508df2569
commit 7718ada4ed
2 changed files with 35 additions and 1 deletions

View File

@@ -6,6 +6,8 @@ import uuid
import comfy.utils
import comfy.model_management
from comfy.types import UnetWrapperFunction
def apply_weight_decompose(dora_scale, weight):
weight_norm = (
@@ -117,7 +119,7 @@ class ModelPatcher:
if disable_cfg1_optimization:
self.model_options["disable_cfg1_optimization"] = True
def set_model_unet_function_wrapper(self, unet_wrapper_function):
def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
self.model_options["model_function_wrapper"] = unet_wrapper_function
def set_model_denoise_mask_function(self, denoise_mask_function):