From e20b911bcc4fe5b98d94e5c4eab1cbd173471dd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C7=9Dunbo=C9=AFi=C9=94s?= Date: Thu, 20 Feb 2025 04:52:27 +0900 Subject: [PATCH] Add files via upload --- Lighting_and_Color.py | 280 +++++++++++++++++++++++++++++++++++++++ __init__.py | 19 +++ nodes.py | 302 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 601 insertions(+) create mode 100644 Lighting_and_Color.py create mode 100644 __init__.py create mode 100644 nodes.py diff --git a/Lighting_and_Color.py b/Lighting_and_Color.py new file mode 100644 index 0000000..bbc823b --- /dev/null +++ b/Lighting_and_Color.py @@ -0,0 +1,280 @@ +from PIL import Image, ImageEnhance, ImageFilter +import numpy as np +import torch +import cv2 as cv + +def tensor2pil(tensor): + """ + PyTorchテンソルをPIL画像に変換する。 + BHWC形式で値が[0, 1]の範囲のテンソルを想定。 + """ + # テンソルがCPUにあることを確認 + tensor = tensor.cpu() + + # numpy配列に変換し、0-255の範囲に正規化 + img_array = tensor.numpy() + img_array = np.clip(img_array * 255.0, 0, 255).astype(np.uint8) + + # バッチ処理の場合は最初の画像を取得 + if len(img_array.shape) == 4: + img_array = img_array[0] # バッチ次元を削除 + + return Image.fromarray(img_array, mode='RGB') + +def pil2tensor(image, original_shape): + """ + PIL画像をPyTorchテンソルに変換し、元の形状に合わせる。 + BHWC形式で値が[0, 1]の範囲のテンソルを返す。 + """ + # numpy配列に変換し、[0, 1]に正規化 + img_array = np.array(image).astype(np.float32) / 255.0 + + # 元の形状に合わせてバッチ次元を追加 + img_array = img_array[np.newaxis, ...] + + # PyTorchテンソルに変換 + tensor = torch.from_numpy(img_array).float() + + return tensor + +def medianFilter(image, radius, num_samples, threshold): + """ + 画像に高品質なメディアンフィルタを適用する + """ + # PILからCV2形式に変換 + cv_image = cv.cvtColor(np.array(image), cv.COLOR_RGB2BGR) + + # エッジを保持しながらスムージングを適用 + blurred = cv.bilateralFilter(cv_image, radius, num_samples, threshold) + + # PIL形式に戻す + return Image.fromarray(cv.cvtColor(blurred, cv.COLOR_BGR2RGB)) + +class FluxLightingAndColor: + """ + FluxLightingAndColor Version 1.2 + 画像の照明と色調を調整するためのノードクラス + + 主な機能: + - 彩度調整 + - 被写界深度(DoF)処理 + - 最適化された処理順序 + - デバッグ出力 + """ + + @classmethod + def INPUT_TYPES(s): + """ + 入力パラメータの定義 + required: 必須パラメータ + - image: 入力画像 + - black/mid/white_level: レベル調整用パラメータ + - red/green/blue_level: 各色チャンネルの強度 + - brightness: 明るさ + - saturation: 彩度 + + optional: オプションパラメータ + - depth: 深度マップ画像 + - dof_mode: 被写界深度エフェクトモード + - dof_radius: ぼかしの半径 + - dof_samples: サンプル数 + - debug_mode: デバッグ出力の有無 + """ + return { + "required": { + "image": ("IMAGE",), + "black_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "mid_level": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "white_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "red_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "green_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "blue_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + }, + "optional": { + "depth": ("IMAGE",), + "dof_mode": (["none", "mock", "gaussian", "box"],), + "dof_radius": ("INT", {"default": 8, "min": 1, "max": 128, "step": 1}), + "dof_samples": ("INT", {"default": 1, "min": 1, "max": 3, "step": 1}), + "debug_mode": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_lighting_and_color" + CATEGORY = "image/adjustments" + + def apply_dof(self, img, depth_map=None, mode='none', radius=8, samples=1, debug=False): + """ + 被写界深度(DoF)エフェクトを適用する + + Parameters: + - img: 元画像 + - depth_map: 深度マップ + - mode: エフェクトモード (none/mock/gaussian/box) + - radius: ぼかしの半径 + - samples: サンプル数 + - debug: デバッグ出力フラグ + """ + if mode == 'none' or depth_map is None: + if debug: + print("DoF: Skipped (mode: none or no depth map)") + return img + + if debug: + print(f"DoF: Applying {mode} blur with radius {radius} and {samples} samples") + + # Resize depth map to match image size and convert to grayscale + depth_map = depth_map.resize(img.size).convert('L') + + # Apply blur based on selected mode + if mode == 'mock': + blurred = medianFilter(img, radius, (radius * 1500), 75) + elif mode == 'gaussian': + blurred = img.filter(ImageFilter.GaussianBlur(radius=radius)) + elif mode == 'box': + blurred = img.filter(ImageFilter.BoxBlur(radius)) + else: + return img + + blurred = blurred.convert(img.mode) + + # Apply multiple samples if requested + if samples > 1: + result = None + for i in range(samples): + if not result: + result = Image.composite(img, blurred, depth_map) + else: + result = Image.composite(result, blurred, depth_map) + if debug: + print(f"DoF: Applied sample {i+1}/{samples}") + else: + result = Image.composite(img, blurred, depth_map).convert('RGB') + + return result + + def adjust_levels(self, img_array, black_level, mid_level, white_level, debug=False): + """ + 画像のレベル調整を行う + + Parameters: + - img_array: 画像配列 + - black_level: 黒レベル + - mid_level: 中間トーン + - white_level: 白レベル + - debug: デバッグ出力フラグ + """ + if debug: + print(f"Levels: Adjusting (black: {black_level}, mid: {mid_level}, white: {white_level})") + # Apply level adjustments + img_array = (img_array - black_level) / (white_level - black_level) + img_array = np.clip(img_array, 0, 1) + return img_array + + def adjust_channels(self, img_array, red_level, green_level, blue_level, debug=False): + """ + RGB各チャンネルの強度を調整する + + Parameters: + - img_array: 画像配列 + - red_level: 赤チャンネルの強度 + - green_level: 緑チャンネルの強度 + - blue_level: 青チャンネルの強度 + - debug: デバッグ出力フラグ + """ + if debug: + print(f"Channels: Adjusting (R: {red_level}, G: {green_level}, B: {blue_level})") + # Split and adjust each channel + img_array[:,:,0] = np.clip(img_array[:,:,0] * red_level, 0, 1) + img_array[:,:,1] = np.clip(img_array[:,:,1] * green_level, 0, 1) + img_array[:,:,2] = np.clip(img_array[:,:,2] * blue_level, 0, 1) + return img_array + + def apply_lighting_and_color(self, image, black_level, mid_level, white_level, + red_level, green_level, blue_level, brightness, saturation, + depth=None, dof_mode="none", dof_radius=8, dof_samples=1, + debug_mode=False): + """ + メインの処理関数。以下の順序で画像処理を実行: + 1. 入力テンソルをPIL画像に変換 + 2. 被写界深度エフェクトの適用（深度マップがある場合） + 3. numpy配列に変換 + 4. 明るさの調整 + 5. レベル調整 + 6. チャンネル調整 + 7. PIL画像に再変換 + 8. 彩度の調整 + 9. 最終的なコントラスト調整 + 10. テンソルに再変換して返却 + """ + try: + if debug_mode: + print("\n=== FluxLightingAndColor v1.2 Starting ===") + print(f"Input tensor shape: {image.shape}, dtype: {image.dtype}") + + # 1. 入力テンソルをPIL画像に変換 + img_pil = tensor2pil(image) + if debug_mode: + print("Step 1: Converted input tensor to PIL image") + + # 2. 深度マップがある場合はDoFを適用 + if depth is not None and dof_mode != "none": + depth_pil = tensor2pil(depth) + img_pil = self.apply_dof(img_pil, depth_pil, dof_mode, dof_radius, dof_samples, debug_mode) + if debug_mode: + print("Step 2: Applied depth of field effect") + + # 3. 処理用にnumpy配列に変換 + img_array = np.array(img_pil).astype(float) / 255.0 + + # 4. トーン調整を適用 + if debug_mode: + print(f"Step 4: Applying brightness boost: {brightness}") + img_array = np.power(img_array, 0.7) * brightness + + # 5. レベル調整を適用 + img_array = self.adjust_levels(img_array, black_level, mid_level, white_level, debug_mode) + + # 6. チャンネル調整を適用 + img_array = self.adjust_channels(img_array, red_level, green_level, blue_level, debug_mode) + + # 7. エンハンス処理用にPIL画像に再変換 + processed = Image.fromarray((np.clip(img_array * 255.0, 0, 255)).astype(np.uint8)) + + # 8. 彩度を適用 + if debug_mode: + print(f"Step 8: Applying saturation: {saturation}") + enhancer = ImageEnhance.Color(processed) + processed = enhancer.enhance(saturation) + + # 9. 最終的なコントラストを適用 + if debug_mode: + print("Step 9: Applying final contrast boost (1.3)") + enhancer = ImageEnhance.Contrast(processed) + processed = enhancer.enhance(1.3) + + # 10. テンソルに再変換 + result = pil2tensor(processed, image.shape) + + if debug_mode: + print(f"Output tensor shape: {result.shape}, dtype: {result.dtype}") + print("=== Processing complete ===\n") + + return (result,) + + except Exception as e: + print(f"Error in apply_lighting_and_color: {str(e)}") + import traceback + traceback.print_exc() + return (image,) + +# Node registration +NODE_CLASS_MAPPINGS = { + "FluxLightingAndColor": FluxLightingAndColor +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxLightingAndColor": "Flux Lighting & Color v1.2" +} \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..456a9b6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,19 @@ +from .nodes import (ControlNetSwitch, ImageSwitch, LatentSwitch, FluxSamplerPuLID, + NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) +from .Lighting_and_Color import FluxLightingAndColor +from .load_input_output_image import NODE_CLASS_MAPPINGS as LOAD_IMAGE_NODES +from .load_input_output_image import NODE_DISPLAY_NAME_MAPPINGS as LOAD_IMAGE_DISPLAY_NAMES + +# ノードマッピングに追加 +NODE_CLASS_MAPPINGS.update({ + "FluxLightingAndColor": FluxLightingAndColor, +}) +NODE_CLASS_MAPPINGS.update(LOAD_IMAGE_NODES) + +# 表示名マッピングに追加 +NODE_DISPLAY_NAME_MAPPINGS.update({ + "FluxLightingAndColor": "Flux Lighting & Color", +}) +NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_IMAGE_DISPLAY_NAMES) + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..e916bc9 --- /dev/null +++ b/nodes.py @@ -0,0 +1,302 @@ +# XLabsのControlNet条件を切り替えるカスタムノード +class ControlNetSwitch: + # ノードの初期化 + def __init__(self): + self.type = "ControlNetSwitch" + print("ControlNetSwitch initialized") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "controlnet_condition_1": ("ControlNetCondition",), + "controlnet_condition_2": ("ControlNetCondition",), + "use_first": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("ControlNetCondition",) + RETURN_NAMES = ("controlnet_condition",) + FUNCTION = "switch" + CATEGORY = "XLabsNodes" + + def switch(self, controlnet_condition_1, controlnet_condition_2, use_first): + print(f"Switching ControlNet conditions. Using {'first' if use_first else 'second'} condition") + return (controlnet_condition_1 if use_first else controlnet_condition_2,) + +class ImageSwitch: + def __init__(self): + self.type = "ImageSwitch" + print("ImageSwitch initialized") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_1": ("IMAGE",), + "image_2": ("IMAGE",), + "use_first": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "switch" + CATEGORY = "image" + + def switch(self, image_1, image_2, use_first): + print(f"Switching images. Using {'first' if use_first else 'second'} image") + return (image_1 if use_first else image_2,) + +class LatentSwitch: + def __init__(self): + self.type = "LatentSwitch" + print("LatentSwitch initialized") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent_1": ("LATENT",), + "latent_2": ("LATENT",), + "use_first": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "switch" + CATEGORY = "latent" + + def switch(self, latent_1, latent_2, use_first): + print(f"Switching latents. Using {'first' if use_first else 'second'} latent") + print(f"Latent 1 shape: {latent_1['samples'].shape}") + print(f"Latent 2 shape: {latent_2['samples'].shape}") + result = latent_1 if use_first else latent_2 + print(f"Result shape: {result['samples'].shape}") + return (result,) + +class FluxSamplerPuLID: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "conditioning": ("CONDITIONING",), + "neg_conditioning": ("CONDITIONING",), + "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 100}), + "timestep_to_start_cfg": ("INT", {"default": 20, "min": 0, "max": 100}), + "true_gs": ("FLOAT", {"default": 3, "min": 0, "max": 100}), + "image_to_image_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "max_shift": ("FLOAT", {"default": 1.15, "min": 0.0, "max": 2.0, "step": 0.01}), + "base_shift": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}), + }, + "optional": { + "latent_image": ("LATENT", {"default": None}), + "controlnet_condition": ("ControlNetCondition", {"default": None}), + } + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "sampling" + CATEGORY = "XLabsNodes" + + def sampling(self, model, conditioning, neg_conditioning, + noise_seed, steps, timestep_to_start_cfg, true_gs, + image_to_image_strength, denoise_strength, + max_shift, base_shift, + latent_image=None, controlnet_condition=None + ): + import torch + import comfy.model_management as mm + from comfy_extras.nodes_model_advanced import ModelSamplingFlux + import latent_preview + import importlib.util + import os + import sys + + # モジュールを動的にインポート + def import_from_path(module_name, file_path): + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module # これを追加 + spec.loader.exec_module(module) + return module + + root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + x_flux_path = os.path.join(root_path, "x-flux-comfyui") + + try: + # 各モジュールを動的にインポート + layers_module = import_from_path("layers", os.path.join(x_flux_path, "layers.py")) + sampling_module = import_from_path("sampling", os.path.join(x_flux_path, "sampling.py")) + utils_module = import_from_path("utils", os.path.join(x_flux_path, "utils.py")) + + # 必要な関数とクラスを取得 + get_noise = sampling_module.get_noise + prepare = sampling_module.prepare + get_schedule = sampling_module.get_schedule + denoise = sampling_module.denoise + denoise_controlnet = sampling_module.denoise_controlnet + unpack = sampling_module.unpack + LATENT_PROCESSOR_COMFY = utils_module.LATENT_PROCESSOR_COMFY + ControlNetContainer = utils_module.ControlNetContainer + DoubleStreamMixerProcessor = layers_module.DoubleStreamMixerProcessor + timestep_embedding = layers_module.timestep_embedding + + except Exception as e: + print(f"Error importing x-flux-comfyui modules: {str(e)}") + print(f"Looking in path: {x_flux_path}") + print(f"Available files: {os.listdir(x_flux_path)}") + raise + + # PuLID Fluxのモデル処理を追加 + modelsamplingflux = ModelSamplingFlux() + width = latent_image["samples"].shape[3]*8 + height = latent_image["samples"].shape[2]*8 + work_model = modelsamplingflux.patch(model, max_shift, base_shift, width, height)[0] + + additional_steps = 11 if controlnet_condition is None else 12 + mm.load_model_gpu(work_model) + inmodel = work_model.model + + try: + guidance = conditioning[0][1]['guidance'] + except: + guidance = 1.0 + + device = mm.get_torch_device() + if torch.backends.mps.is_available(): + device = torch.device("mps") + if torch.cuda.is_bf16_supported(): + dtype_model = torch.bfloat16 + else: + dtype_model = torch.float16 + + offload_device = mm.unet_offload_device() + torch.manual_seed(noise_seed) + + bc, c, h, w = latent_image['samples'].shape + height = (h//2) * 16 + width = (w//2) * 16 + + x = get_noise( + bc, height, width, device=device, + dtype=dtype_model, seed=noise_seed + ) + orig_x = None + if c==16: + orig_x = latent_image['samples'] + lat_processor2 = LATENT_PROCESSOR_COMFY() + orig_x = lat_processor2.go_back(orig_x) + orig_x = orig_x.to(device, dtype=dtype_model) + + timesteps = get_schedule( + steps, + (width // 8) * (height // 8) // 4, + shift=True, + ) + try: + inmodel.to(device) + except: + pass + x.to(device) + + inmodel.diffusion_model.to(device) + inp_cond = prepare(conditioning[0][0], conditioning[0][1]['pooled_output'], img=x) + neg_inp_cond = prepare(neg_conditioning[0][0], neg_conditioning[0][1]['pooled_output'], img=x) + + if denoise_strength <= 0.99: + try: + timesteps = timesteps[:int(len(timesteps)*denoise_strength)] + except: + pass + + x0_output = {} + callback = latent_preview.prepare_callback(model, len(timesteps) - 1, x0_output) + + if controlnet_condition is None: + x = denoise( + inmodel.diffusion_model, **inp_cond, timesteps=timesteps, guidance=guidance, + timestep_to_start_cfg=timestep_to_start_cfg, + neg_txt=neg_inp_cond['txt'], + neg_txt_ids=neg_inp_cond['txt_ids'], + neg_vec=neg_inp_cond['vec'], + true_gs=true_gs, + image2image_strength=image_to_image_strength, + orig_image=orig_x, + callback=callback, + width=width, + height=height, + ) + else: + def prepare_controlnet_condition(controlnet_condition): + controlnet = controlnet_condition['model'] + controlnet_image = controlnet_condition['img'] + controlnet_image = torch.nn.functional.interpolate( + controlnet_image, size=(height, width), scale_factor=None, mode='bicubic',) + controlnet_strength = controlnet_condition['controlnet_strength'] + controlnet_start = controlnet_condition['start'] + controlnet_end = controlnet_condition['end'] + controlnet.to(device, dtype=dtype_model) + controlnet_image = controlnet_image.to(device, dtype=dtype_model) + return { + "img": controlnet_image, + "controlnet_strength": controlnet_strength, + "model": controlnet, + "start": controlnet_start, + "end": controlnet_end, + } + + cnet_conditions = [prepare_controlnet_condition(el) for el in controlnet_condition] + containers = [] + for el in cnet_conditions: + start_step = int(el['start']*len(timesteps)) + end_step = int(el['end']*len(timesteps)) + container = ControlNetContainer(el['model'], el['img'], el['controlnet_strength'], start_step, end_step) + containers.append(container) + + mm.load_models_gpu([work_model,]) + + total_steps = len(timesteps) + + x = denoise_controlnet( + inmodel.diffusion_model, **inp_cond, + controlnets_container=containers, + timesteps=timesteps, guidance=guidance, + timestep_to_start_cfg=timestep_to_start_cfg, + neg_txt=neg_inp_cond['txt'], + neg_txt_ids=neg_inp_cond['txt_ids'], + neg_vec=neg_inp_cond['vec'], + true_gs=true_gs, + image2image_strength=image_to_image_strength, + orig_image=orig_x, + callback=callback, + width=width, + height=height, + ) + + x = unpack(x, height, width) + lat_processor = LATENT_PROCESSOR_COMFY() + x = lat_processor(x) + lat_ret = {"samples": x} + + return (lat_ret,) + +# ノードの登録 +NODE_CLASS_MAPPINGS = { + "ControlNetSwitch": ControlNetSwitch, + "ImageSwitch": ImageSwitch, + "LatentSwitch": LatentSwitch, + "FluxSamplerPuLID": FluxSamplerPuLID +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ControlNetSwitch": "ControlNet Switcher", + "ImageSwitch": "Image Switcher", + "LatentSwitch": "Latent Switcher", + "FluxSamplerPuLID": "Flux Sampler For PuLID" +} \ No newline at end of file