import os import random import sys import logging import torch from typing import Union, Sequence, Mapping, Any import shutil # 设置日志 logger = logging.getLogger(__name__) # 从索引获取值的辅助函数 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: try: return obj[index] except KeyError: return obj["result"][index] # 获取ComfyUI根目录 def get_comfyui_dir(): # 从当前目录向上一级查找ComfyUI根目录 current_dir = os.path.dirname(os.path.abspath(__file__)) comfyui_dir = os.path.dirname(current_dir) # 当前目录的上一级目录 return comfyui_dir # 图像填充服务函数 def generate_fill_image( prompt: str, prompt_negative: str = "", input_image_path: str = None, style_image: str = None, reference_strength: float = 0.8, guidance: float = 30.0, steps: int = 28, seed: int = -1 ): """ 使用FLUX Fill模型进行图像填充，可选择使用参考图像 参数: prompt (str): 正面提示词 prompt_negative (str): 负面提示词 input_image_path (str): 输入图像路径（带有透明区域作为填充遮罩） style_image (str): 参考图像路径（用于指导填充风格和内容） reference_strength (float): 参考图像的影响强度，0-1之间，默认0.8 guidance (float): 引导强度，默认30.0 steps (int): 采样步数，默认28 seed (int): 随机种子，-1表示随机生成 返回: str: 生成图像的保存路径 """ logger.info("开始图像填充过程...") # 使用随机种子（如果指定为-1） if seed == -1: seed = random.randint(1, 2**64) # 导入必要组件 from nodes import ( NODE_CLASS_MAPPINGS, SaveImage, VAELoader, UNETLoader, DualCLIPLoader, CLIPTextEncode, VAEDecode, LoadImage, InpaintModelConditioning, CLIPVisionEncode, ConditioningCombine, CLIPVisionLoader,StyleModelLoader,StyleModelApply ) logger.info("加载模型...") # 加载CLIP模型 dualcliploader = DualCLIPLoader() CLIP_MODEL = dualcliploader.load_clip( clip_name1="clip_l.safetensors", clip_name2="t5/t5xxl_fp16.safetensors", type="flux", device="default" ) # 加载CLIP Vision模型 clipvisionloader = CLIPVisionLoader() CLIP_VISION_MODEL = clipvisionloader.load_clip( clip_name="sigclip_vision_patch14_384.safetensors" ) # 加载VAE模型 vaeloader = VAELoader() VAE_MODEL = vaeloader.load_vae(vae_name="FLUX1/ae.safetensors") # 加载UNET模型（Fill模型） unetloader = UNETLoader() UNET_MODEL = unetloader.load_unet( unet_name="flux1-fill-dev.safetensors", weight_dtype="default" ) # 初始化其他节点 cliptextencode = CLIPTextEncode() clipvisionencode = NODE_CLASS_MAPPINGS["CLIPVisionEncode"]() conditioningcombine = NODE_CLASS_MAPPINGS["ConditioningCombine"]() fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]() differentialdiffusion = NODE_CLASS_MAPPINGS["DifferentialDiffusion"]() ksampler = NODE_CLASS_MAPPINGS["KSampler"]() vaedecode = VAEDecode() saveimage = SaveImage() loadimage = LoadImage() inpaintmodelconditioning = InpaintModelConditioning() # 把输入图像复制到ComfyUI的input目录 comfyui_dir = get_comfyui_dir() input_dir = os.path.join(comfyui_dir, "input") os.makedirs(input_dir, exist_ok=True) # 创建临时文件名以避免冲突 input_filename = f"input_{random.randint(1000, 9999)}.png" ref_filename = None if style_image: ref_filename = f"ref_{random.randint(1000, 9999)}.png" input_comfy_path = os.path.join(input_dir, input_filename) shutil.copy(input_image_path, input_comfy_path) ref_comfy_path = None if style_image: ref_comfy_path = os.path.join(input_dir, ref_filename) shutil.copy(style_image, ref_comfy_path) # 预加载模型到GPU from comfy import model_management model_loaders = [CLIP_MODEL, VAE_MODEL, UNET_MODEL, CLIP_VISION_MODEL] model_management.load_models_gpu([ loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders ]) # 图像生成流程 with torch.inference_mode(): # 文本编码 text_encoded = cliptextencode.encode( text=prompt, clip=get_value_at_index(CLIP_MODEL, 0), ) negative_encoded = cliptextencode.encode( text=prompt_negative, clip=get_value_at_index(CLIP_MODEL, 0), ) # 应用FluxGuidance flux_guided = fluxguidance.append( guidance=guidance, conditioning=get_value_at_index(text_encoded, 0), ) # 如果有参考图像，使用StyleModelApplyAdvanced处理 final_positive_cond = get_value_at_index(flux_guided, 0) if style_image: try: # 初始化StyleModelLoader和StyleModelApplyAdvanced节点 stylemodelloader = StyleModelLoader() STYLE_MODEL = stylemodelloader.load_style_model( style_model_name="flux1-redux-dev.safetensors" ) stylemodelapplyadvanced = NODE_CLASS_MAPPINGS["StyleModelApplyAdvanced"]() # 应用风格 # 加载参考图像 ref_img = loadimage.load_image(image=ref_filename) # 使用CLIP Vision编码参考图像 vision_encoded = clipvisionencode.encode( crop="center", # 中心裁剪 clip_vision=get_value_at_index(CLIP_VISION_MODEL, 0), image=get_value_at_index(ref_img, 0) ) # 应用风格 - 将参考图像特征应用到条件上 style_applied = stylemodelapplyadvanced.apply_stylemodel( conditioning=final_positive_cond, # 条件 style_model=get_value_at_index(STYLE_MODEL, 0), # 风格模型 clip_vision_output=get_value_at_index(vision_encoded, 0), # 参考图像特征 strength=0.5, # 风格强度参数 ) # 获取应用风格后的条件 final_positive_cond = get_value_at_index(style_applied, 0) logger.info("成功应用参考图像风格") except Exception as e: logger.error(f"应用参考图像风格时出错: {e}") # 继续使用原始条件，不中断处理 logger.info("继续使用原始条件而不应用参考图像") # 加载图像（同时获取图像和遮罩） loaded_image = loadimage.load_image(image=input_filename) input_img = get_value_at_index(loaded_image, 0) # 第一个输出是图像 mask_img = get_value_at_index(loaded_image, 1) # 第二个输出是遮罩 # 应用DifferentialDiffusion diff_model = differentialdiffusion.apply( model=get_value_at_index(UNET_MODEL, 0), ) # 准备InpaintModelConditioning inpaint_cond = inpaintmodelconditioning.encode( positive=final_positive_cond, negative=get_value_at_index(negative_encoded, 0), vae=get_value_at_index(VAE_MODEL, 0), pixels=input_img, mask=mask_img, noise_mask=False, ) # 从inpaint_cond获取正面条件、负面条件和潜在图像 positive_cond = get_value_at_index(inpaint_cond, 0) negative_cond = get_value_at_index(inpaint_cond, 1) latent_image = get_value_at_index(inpaint_cond, 2) # 采样 sampled = ksampler.sample( model=get_value_at_index(diff_model, 0), positive=positive_cond, negative=negative_cond, latent_image=latent_image, seed=seed, steps=steps, cfg=1.0, # 使用FluxGuidance时，cfg通常设为1.0 sampler_name="euler", scheduler="normal", denoise=1.0, ) # 解码 decoded = vaedecode.decode( samples=get_value_at_index(sampled, 0), vae=get_value_at_index(VAE_MODEL, 0), ) # 设置保存的文件名前缀 prefix = "Flux_Fill" # 保存图像 saved = saveimage.save_images( filename_prefix=prefix, images=get_value_at_index(decoded, 0), ) # 获取保存的图像路径 saved_path = os.path.join(comfyui_dir, "output", saved['ui']['images'][0]['filename']) logger.info(f"图像保存到 {saved_path}") # 清理临时文件 try: os.remove(input_comfy_path) if ref_comfy_path: os.remove(ref_comfy_path) except Exception as e: logger.warning(f"清理临时文件时出错: {e}") return saved_path