nodes/MiddleKD-ComfyUI-productfix
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix: add node class variables

Browse Source
This commit is contained in:
ml-favorfit-develop
2024-07-05 09:28:20 +09:00
parent 81bada6233
commit 665f453dcc
2 changed files with 126 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
advanced_sampler.py
+4 -4
View File
@@ -6,7 +6,7 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
@wraps(original_ksampler_call_fn)
def wrapper(self, x, sigma, denoise_mask, model_options={}, seed=None):
# latent 주입 옵션 확인
# latent injection 옵션 확인
latent_inject_options = model_options.get("is_latent_inject")
if latent_inject_options is not None:
start_sigma = latent_inject_options["start_sigma"]
@@ -17,7 +17,7 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
if "denoise_mask_function" in model_options:
denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
# 이미지 주입 로직
# latent injection 로직
if latent_inject_options is not None:
if sigma.item() > end_sigma and sigma.item() < start_sigma:
latent_mask = 1. - denoise_mask
@@ -31,7 +31,7 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
# 모델 실행
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
# 후처리 (이미지 주입이 아닌 경우에만)
# 후처리 (latent injection이 아닌 경우에만)
if denoise_mask is not None:
if latent_inject_options is not None:
pass
@@ -41,7 +41,7 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
# 원래 함수로 복구
if self.sigmas[-2].item() >= sigma.item():
KSamplerX0Inpaint.__call__ = original_ksampler_call_fn
logging.debug("\033[94m[middlek image injection] KSamplerX0Inpaint.__call__ return to original_ksampler_call_fn\033[0m")
logging.debug("\033[94m[middlek latent injection] KSamplerX0Inpaint.__call__ return to original_ksampler_call_fn\033[0m")
return out
node.py
+122 -11
View File
@@ -13,7 +13,13 @@ from .utils import (simple_resize, add_detail_transfer)
# ModelPatcher의 calculate_weight 메서드를 초기화하는 클래스
class ResetModelPatcherCalculateWeight:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(s):
return {"required": {"model":("MODEL", ),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "reset_moodelpatcher_weight"
CATEGORY = "productfix"
def reset_moodelpatcher_weight(self, model:ModelPatcher):
# 원래의 calculate_weight 메서드로 복원
@@ -24,18 +30,45 @@ class ResetModelPatcherCalculateWeight:
# 잠재 공간에 이미지를 주입하는 클래스
class ApplyLatentInjection:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(s):
return {"required": {"model":("MODEL", ),
"latents": ("LATENT", ),
"inject_image_embed": ("LATENT", ),
"inject_mask": ("MASK",),
"start_sigma": ("FLOAT", {"default": 15.0}),
"end_sigma": ("FLOAT", {"default": 0.0})
}}
RETURN_TYPES = ("MODEL", "LATENT",)
FUNCTION = "apply_latent_injection"
CATEGORY = "productfix"
def apply_latent_injection(self, model, latents, inject_image_embed, inject_mask, start_sigma, end_sigma):
device = model_management.get_torch_device()
dtype = model_management.VAE_DTYPES[0]
# KSamplerX0Inpaint의 __call__ 메서드를 수정된 버전으로 교체
original_ksampler_call_fn = KSamplerX0Inpaint.__call__
KSamplerX0Inpaint.__call__ = inject_ksamplerx0inpaint_call(original_ksampler_call_fn)
logging.info("\033[94m[middlek image injection] KSamplerX0Inpaint.__call__ is injected to inject_ksamplerx0inpaint_call\033[0m")
logging.info("\033[94m[middlek latent injection] KSamplerX0Inpaint.__call__ is injected to inject_ksamplerx0inpaint_call\033[0m")
# 잠재 공간 및 마스크 준비
# ... (데이터 형식 변환 및 디바이스 이동)
# (데이터 형식 변환 및 디바이스 이동)
if isinstance(inject_image_embed, dict):
inject_image_embed = inject_image_embed["samples"]
inject_image_embed = inject_image_embed.to(dtype=dtype)
b, c, h, w = inject_image_embed.shape
if len(inject_mask.shape) != 4:
inject_mask = inject_mask.unsqueeze(0)
inject_image_embed = inject_image_embed.to(device=device, dtype=dtype)
inject_mask = simple_resize(inject_mask, h, w).to(device=device, dtype=dtype)
latents["samples"] = inject_image_embed
latents["noise_mask"] = inject_mask
# 모델 옵션에 이미지 주입 파라미터 추가
# 모델 옵션에 latent injection 파라미터 추가
if hasattr(model, "model_options"):
model.model_options["is_latent_inject"] = {"start_sigma":start_sigma, "end_sigma":end_sigma}
@@ -43,7 +76,18 @@ class ApplyLatentInjection:
# VQ 모델을 로드하는 클래스
class VQLoader:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"vq_name": ([os.path.basename(cur) for cur in folder_paths.get_filename_list("vae_approx") if "vq" in cur], ),
}
}
RETURN_TYPES = ("VQ", )
FUNCTION = "load_vq"
CATEGORY = "productfix"
def load_vq(self, vq_name):
ckpt = os.path.join(folder_paths.get_folder_paths("vae_approx")[0], vq_name)
@@ -52,7 +96,18 @@ class VQLoader:
# VQ 모델을 사용하여 이미지를 인코딩하는 클래스
class VQEncoder:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"vq": ("VQ",),
"images": ("IMAGE", ),
}
}
RETURN_TYPES = ("LATENT", )
FUNCTION = "encode"
CATEGORY = "productfix"
def encode(self, images, vq):
latents = vqmodel_encode(images, vq)
@@ -61,7 +116,18 @@ class VQEncoder:
# VQ 모델을 사용하여 잠재 공간을 디코딩하는 클래스
class VQDecoder:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"vq": ("VQ",),
"latents": ("LATENT", ),
}
}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "decode"
CATEGORY = "productfix"
def decode(self, latents, vq):
if isinstance(latents, dict):
@@ -71,7 +137,22 @@ class VQDecoder:
# 이미지에서 텍스트 마스크를 생성하는 클래스
class GetTextMask:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(s):
return {"required": {
"image": ("IMAGE",),
"languages": (
get_languages() + ["not use"],
{"default": "not use"},
),
"codes": (
"STRING",
{"default": "en,ko", "multiline": False},
),
}}
RETURN_TYPES = ("MASK",)
FUNCTION = "get_text_mask"
CATEGORY = "productfix"
def get_text_mask(self, image, languages:str, codes:str):
# 언어 설정에 따라 타겟 언어 결정
@@ -85,7 +166,22 @@ class GetTextMask:
# 디테일 전송을 수행하는 클래스 (이미지 도메인)
class DetailTransferAdd:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(s):
return {"required": {
"target": ("IMAGE", ),
"source": ("IMAGE", ),
"blur": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.01}),
"blend_ratio": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001, "round": 0.001}),
},
"optional": {
"mask": ("MASK", ),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "detail_transfer_add"
CATEGORY = "productfix"
def detail_transfer_add(self, target, source, blur, blend_ratio, mask=None):
output_image = add_detail_transfer(target, source, blur, blend_ratio, mask)
@@ -93,7 +189,22 @@ class DetailTransferAdd:
# 디테일 전송을 수행하는 클래스 (잠재 공간 도메인)
class DetailTransferLatentAdd:
# ... (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY 정의)
@classmethod
def INPUT_TYPES(s):
return {"required": {
"target": ("LATENT", ),
"source": ("LATENT", ),
"blur": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.01}),
"blend_ratio": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001, "round": 0.001}),
},
"optional": {
"mask": ("MASK", ),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "detail_transfer_add"
CATEGORY = "productfix"
def detail_transfer_add(self, target, source, blur, blend_ratio, mask=None):
# 잠재 공간을 이미지 형식으로 변환