ai-swap/backend/modules/comfyui/ai_swap.py

import json
import random
import websocket
from PIL import Image
from datetime import datetime
from backend.modules.comfyui.image_tools import upload_image, show_gif
from backend.modules.comfyui.network_tools import get_images
from backend.utils.logger_config import setup_logger
from backend.utils.system_config import Config
from backend.config.prompt_config import FACE_PROMPT, CLOTH_PROMPT, PROMPT_PROMPT

system_config = Config('./backend/config/workflow_config.json')

logger = setup_logger(__name__)

def parse_workflow(prompt, face_img, cloth_img, config, face_prompt, cloth_prompt, prompt_prompt) -> dict:
    try:
        # 上传图片
        logger.info("开始上传人脸图片...")
        input_face_img = upload_image(face_img, config)
        logger.info("人脸图片上传成功")
        
        logger.info("开始上传服装图片...")
        input_cloth_img = upload_image(cloth_img, config)
        logger.info("服装图片上传成功")
        
        logger.info(f'Work flow args config: workflow-{config.workflowfile}, prompt-{prompt}, seed-{config.seed}')
        
        # 加载工作流配置
        try:
            with open(config.workflowfile, "r", encoding="utf-8") as workflow:
                workflow_config = json.load(workflow)
        except FileNotFoundError:
            logger.error(f"工作流配置文件不存在: {config.workflowfile}")
            raise RuntimeError(f"工作流配置文件不存在: {config.workflowfile}")
        except json.JSONDecodeError as e:
            logger.error(f"工作流配置文件格式错误: {str(e)}")
            raise RuntimeError(f"工作流配置文件格式错误: {str(e)}")
        
        # 更新工作流参数（旧版）
        # workflow_config["268"]["inputs"]["text"] = prompt
        # workflow_config["271"]["inputs"]["image"] = input_face_img
        # workflow_config["265"]["inputs"]["image"] = input_cloth_img
        # workflow_config["267"]["inputs"]["seed"] = config.seed

        # 更新工作流参数（1014版）
        workflow_config["251"]["inputs"]["prompt"] = prompt
        workflow_config["248"]["inputs"]["image"] = input_face_img
        workflow_config["249"]["inputs"]["image"] = input_cloth_img
        workflow_config["251"]["inputs"]["seed"] = config.seed

        # 更新工作流参数（1028版）
        # workflow_config["253"]["inputs"]["prompt"] = prompt
        # workflow_config["248"]["inputs"]["image"] = input_face_img
        # workflow_config["249"]["inputs"]["image"] = input_cloth_img
        
        # TODO: 开放人脸识别、服装识别、提示词优化 三类系统提示词
        logger.info(f"system_prompt unable: {face_prompt}, {cloth_prompt}, {prompt_prompt}")
        # workflow_config["217"]["inputs"]["system_prompt"] = face_prompt
        # workflow_config["174"]["inputs"]["system_prompt"] = cloth_prompt
        # workflow_config["231"]["inputs"]["system_prompt"] = prompt_prompt
        
        logger.info("开始执行工作流...")
        return get_images(workflow_config, config)
        
    except Exception as e:
        logger.error(f"工作流解析失败: {str(e)}")
        raise RuntimeError(f"工作流解析失败: {str(e)}")

def ai_swap_process(prompt, face_img, cloth_img, face_prompt=FACE_PROMPT, cloth_prompt=CLOTH_PROMPT, prompt_prompt=PROMPT_PROMPT.format(history="暂无历史记录")) -> list:
    ws = websocket.WebSocket()
    host_url = f"ws://{system_config.server_address}/ws?clientID={system_config.client_id}"
    
    try:
        ws.connect(host_url)
        logger.info(f'Request to: {host_url}')
    except Exception as e:
        logger.error(f"无法连接到ComfyUI服务器: {host_url}, 错误: {str(e)}")
        raise RuntimeError(f"ComfyUI服务器连接失败: {str(e)}")

    # 从kwargs中提取参数
    try:
        images, history_prompt = parse_workflow(prompt, face_img, cloth_img, system_config, face_prompt, cloth_prompt, prompt_prompt)
        logger.info(f"工作流解析完成，获取到 {len(images)} 个节点输出")
    except Exception as e:
        logger.error(f"工作流解析失败: {str(e)}")
        raise RuntimeError(f"工作流解析失败: {str(e)}")
    
    images_cc = []

    # 检查是否有图片输出
    if not images:
        logger.error("没有获取到任何图片输出，请检查ComfyUI服务器状态和工作流配置")
        raise RuntimeError("AI处理未返回任何结果，请检查服务器状态")

    for node_id in images:
        logger.info(f"处理节点 {node_id}，包含 {len(images[node_id])} 张图片")
        for image_data in images[node_id]:
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            GIF_LOCATION = f"{system_config.output_dir}/{system_config.idx}_{system_config.seed}_{timestamp}.png"
            try:
                with open(GIF_LOCATION, "wb") as binary_file:
                    binary_file.write(image_data)
                    show_gif(GIF_LOCATION)
                system_config.idx += 1
                images_cc.append(Image.open(GIF_LOCATION))
                logger.info(f"图片保存成功: {GIF_LOCATION}")
            except Exception as e:
                logger.error(f"图片保存失败: {str(e)}")
                raise RuntimeError(f"图片保存失败: {str(e)}")
    
    logger.info(f'Prompt queue finished! The return result: {len(images_cc)} images')
    
    # 检查是否有处理结果
    if not images_cc:
        logger.error("没有生成任何图片结果")
        raise RuntimeError("AI处理未生成任何图片结果，请检查输入参数和服务器状态")
    
    return images_cc[0], history_prompt

if __name__ == "__main__":

    import numpy as np
    face_img = np.array(Image.open('backend/data/face.png'))
    cloth_img = np.array(Image.open('backend/data/cloth.png'))
    ai_swap_process('美女站在海边', face_img, cloth_img)