qqbot/ai_pic.py

import requests
import json
import time
import random
from websocket import WebSocketApp
from config import config_get
import os

# 全局变量在这里改写：
file_path = config_get('file_path')
ip = config_get("comfyUI_ip")
pro_bad = config_get("pro_bad")

def AI_paint(prompt):
    """
    根据传入的prompt参数获取image_url
    """
    png_path = 'output.png'
    url = "https://api.siliconflow.cn/v1/images/generations"
    web_url = "https://e423-34-68-183-106.ngrok-free.app"
    payload = {
        "model": "black-forest-labs/FLUX.1-schnell",
        "image_size": "1024x1024",
        "prompt": prompt,  # 正面提示词
        "negative_prompt": pro_bad,
        "batch_size": 2,
        "num_inference_steps": 35,
        "guidance_scale": 7.5,

    }

    headers = {
        "Authorization": "Bearer sk-hrascmreejxkirpsdprupkjtqbxtywqbyntzjkdjuficbtok",
        "Content-Type": "application/json"
    }

    response = requests.post(url, json=payload, headers=headers)

    r = response.json()
    image_url = r['images'][0]['url']
    print(image_url)
    return image_url

def generate_instance_id():
    """
    根据当前时间生成图片生成唯一实例id
    """
    current_time = time.strftime("%Y%m%d%H%M%S", time.localtime())
    instance_id = f"{current_time}-gxx12138-key"
    return instance_id

def start_paint(id,data):
    """
    绘图任务下发程序，传入实例id和提示词开始任务
    传入修改好的json格式的data数据，发送到confyUI下发任务
    """
    # 构建符合官方文档格式的请求数据
    request_data = {
        "client_id": id,
        "prompt": data
    }

    # 发送POST请求到ComfyUI的/prompt端点
    url = "http://"+ip+"/prompt"
    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, headers=headers, json=request_data)

    # 检查响应状态码是否为200（表示成功）
    if response.status_code == 200:
        print("任务下发成功")
        print(response.json())
        print("任务实例id为"+id)
    else:
        print(f"请求失败，状态码: {response.status_code}")

def get_websocket_data(client_id):
    """
    获取生图任务运行状态，根据实例id进行跟踪
    """
    output_image_filenames = []
    def on_message(ws, message):
        nonlocal output_image_filenames
        message_dict = json.loads(message)
        if message_dict["type"] == "executed" and message_dict["data"]["node"] == "16":
            output_data = message_dict["data"]
            if "output" in output_data and "images" in output_data["output"]:
                image_list = output_data["output"]["images"]
                output_image_filenames = [image["filename"] for image in image_list]
            ws.close()

        message_dict = eval(message)
        if message_dict["type"] != "crystools.monitor":
            print(f"\r收到消息: {message}",flush=True,end="\n")

    def on_error(ws, error):
        print(f"\r发生错误: {error}",flush=True,end="\n")

    def on_close(ws, close_status_code, close_msg):
        print(f"\r连接已关闭，状态码: {close_status_code}，消息: {close_msg}",flush=True,end="\n")

    def on_open(ws):
        print("\r连接已建立",flush=True,end="\n")

    websocket_url = "ws://"+ip+"/ws?clientId="+client_id
    ws = WebSocketApp(websocket_url,
                      on_open=on_open,
                      on_message=on_message,
                      on_error=on_error,
                      on_close=on_close)
    ws.run_forever()

    return output_image_filenames

def get_pic(id):
    """
    输入prompt_id来查询图片
    """
    url = "http://"+ip+"/history/"+id

    try:
        response = requests.get(url)
        if response.status_code == 200:
            print("请求成功，响应内容：")
            print(response.json())
        else:
            print(f"请求失败，状态码：{response.status_code}")
    except requests.RequestException as e:
        print(f"请求发生错误：{e}")

def get_images_from_filenames(image_filenames):
    """
    获取图片存入output,传入为图片名称列表
    ip和存储地址采用硬编码，如果comfyUI配置不同需要调整ip和路径
    """
    base_url = "http://"+ip+"/view"
    output_folder = f"{file_path}/output"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for filename in image_filenames:
        params = {
            "filename": filename,
            "type": "output"
        }
        response = requests.get(base_url, params=params)
        if response.status_code == 200:
            # 这里可以根据实际需求对获取到的图片内容进行进一步处理，比如保存到本地等
            print(f"成功获取图片 {filename} 的内容")
            # 保存图片到output文件夹下
            output_path = os.path.join(output_folder, filename)
            with open(output_path, 'wb') as f:
                f.write(response.content)
            print(f"成功将图片 {filename} 保存到 {output_path}")
            return output_path
        else:
            print(f"获取图片 {filename} 时出错，状态码: {response.status_code}")

def story_start(prompt):
    """
    下发任务开关，在这里动态处理任务相关参数
    """
    with open('maomao_api.json', 'r', encoding='utf-8') as file:
        data = json.load(file)
    data['3']['inputs']["text"] = prompt
    if str(prompt).startswith("nsfw"):
        data['4']['inputs']["text"] = config_get("pro_bad_sese")
        print("true")
    data["21"]["inputs"]["seed"] = random.randint(1, 2 ** 32 - 1)
    data["5"]["inputs"]["noise_seed"] = random.randint(0, 18446744073709551615)
    print(prompt)
    id = generate_instance_id()
    start_paint(id, data)
    return get_images_from_filenames(get_websocket_data(id))

def story_start_p2p(qq):
    """
    调用p2p的任务下发开关，传入qq号，自动根据user文件夹下的用户名搜索头像来变身
    用户头像存储路径默认user文件夹
    """
    with open('p2p.json', 'r', encoding='utf-8') as file:
        data = json.load(file)
    data['49']['inputs']["directory"] = f"{file_path}/user/{qq}/"
    data["43"]["inputs"]["seed"] = random.randint(1, 2 ** 32 - 1)
    id = generate_instance_id()
    start_paint(id, data)
    return get_images_from_filenames(get_websocket_data(id))

def story_start_p2p_sese(qq):
    """
    调用p2p涩图版的任务下发开关，传入qq号，自动根据user文件夹下的用户名搜索头像来变身
    用户头像存储路径默认user文件夹
    """
    with open('p2p.json', 'r', encoding='utf-8') as file:
        data = json.load(file)
    data['49']['inputs']["directory"] = f"{file_path}/user/{qq}/"
    data["43"]["inputs"]["seed"] = random.randint(1, 2 ** 32 - 1)
    data['3']['inputs']["text"] = "nsfw,1girl"
    data['4']['inputs']["text"] = config_get("pro_bad_sese")
    id = generate_instance_id()
    start_paint(id, data)
    return get_images_from_filenames(get_websocket_data(id))