100 lines
4.1 KiB
Python
100 lines
4.1 KiB
Python
import sys
|
|
from .utils import dynamic_resize
|
|
from denoise_mask_scheduler import DenoiseMaskScheduler
|
|
from advanced_sampler import KSamplerX0Inpaint, inject_ksamplerx0inpaint_call
|
|
import logging
|
|
|
|
def inject_ksampler():
|
|
# KSamplerX0Inpaint의 __call__ 메서드를 수정된 버전으로 교체
|
|
original_ksampler_call_fn = KSamplerX0Inpaint.__call__
|
|
KSamplerX0Inpaint.__call__ = inject_ksamplerx0inpaint_call(original_ksampler_call_fn)
|
|
logging.info("\033[94m[Apply denoise-mask-scheduler] KSamplerX0Inpaint.__call__ is injected to inject_ksamplerx0inpaint_call\033[0m")
|
|
|
|
class ApplyDenoiseMaskSchedulerWithSigma:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"model":("MODEL", ),
|
|
"start_sigma": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 15.0, "step": 0.05}),
|
|
"end_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 15.0, "step": 0.05}),
|
|
"mask_alpha": ("FLOAT", {"default":1.0, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
"scheduler_type": (["skip"], {"default":"skip"}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "denoise-mask-scheduler"
|
|
|
|
def run(self, model, start_sigma, end_sigma, mask_alpha, scheduler_type):
|
|
if scheduler_type == "skip":
|
|
scheduler_type = "skip_with_sigma"
|
|
else:
|
|
raise ValueError(f"{scheduler_type} is not supported for denoise mask scheduler.")
|
|
|
|
denoise_mask_scheduler = DenoiseMaskScheduler(
|
|
scheduler_type=scheduler_type,
|
|
start_sigma=start_sigma,
|
|
end_sigma=end_sigma,
|
|
mask_alpha=mask_alpha,
|
|
)
|
|
|
|
if hasattr(model, "model_options"):
|
|
model.model_options["denoise_mask_function"] = denoise_mask_scheduler.get_denoise_mask
|
|
else:
|
|
raise TypeError(f"'{type(model)}' is not a ComfyUI ModelPatcher type.")
|
|
|
|
inject_ksampler()
|
|
|
|
return (model, )
|
|
|
|
class ApplyDenoiseMaskSchedulerWithStep:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"model":("MODEL", ),
|
|
"start_step": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
|
"end_step": ("INT", {"default": 20, "min": 0, "max": 100, "step": 1}),
|
|
"mask_alpha": ("FLOAT", {"default":1.0, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
"scheduler_type": (["skip"], {"default":"skip"}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "denoise-mask-scheduler"
|
|
|
|
def run(self, model, start_step, end_step, mask_alpha, scheduler_type):
|
|
if scheduler_type == "skip":
|
|
scheduler_type = "skip_with_step"
|
|
else:
|
|
raise ValueError(f"{scheduler_type} is not supported for denoise mask scheduler.")
|
|
|
|
denoise_mask_scheduler = DenoiseMaskScheduler(
|
|
scheduler_type=scheduler_type,
|
|
start_step=start_step,
|
|
end_step=end_step,
|
|
mask_alpha=mask_alpha,
|
|
)
|
|
|
|
if hasattr(model, "model_options"):
|
|
model.model_options["denoise_mask_function"] = denoise_mask_scheduler.get_denoise_mask
|
|
else:
|
|
raise TypeError(f"'{type(model)}' is not a ComfyUI ModelPatcher type.")
|
|
|
|
inject_ksampler()
|
|
|
|
return (model, )
|
|
|
|
# 동적 이미지 리사이즈
|
|
class DynamicImageResize:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"max_pixels": ("INT", {"default": 1024*1024, "min": 500, "max": sys.maxsize, "step": 1}),
|
|
"min_pixels": ("INT", {"default": 512*512, "min": 500, "max": sys.maxsize, "step": 1}),
|
|
}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "dynamic_image_resize"
|
|
|
|
def dynamic_image_resize(self, image, max_pixels:int, min_pixels:int):
|
|
resized_image = dynamic_resize(image, max_pixels=max_pixels, min_pixels=min_pixels)
|
|
return (resized_image, )
|