SD&comfYui網路小說生產器

不上火星不改名發表於2024-03-10
UI
import PySimpleGUI as sg
import os
import glob
from urllib import request, parse
import json
import random
import ctypes

# 假設這是你的OpenAI客戶端配置
from openai import OpenAI

client = OpenAI(
    base_url="https://oneapi.xty.app/v1",
    api_key="sk-JBQJmyrMcAAJJtWb00816e05Ff80438389BaE80834Ea5480"  # 請替換為你的API金鑰
)

# 定義系統指令
instruction = (
    "你是stable diffusion的提示詞生成器，我會給你一個網路文學，如:{0}。請根據這個文學的畫面產生提示語。"
    "請你先對網路文學進行判斷，是古代還是現代。如果是古代，生成一個nijii國風言情風格的提示語。"
    "如果是現代，生成一個nijii中國現代都市網紅風格的提示語，提示語儘量要長，回答不要加引號。"
    "風格偏向：中國風，武俠，美少女，美少年，niji，白瘦幼"
    "只要返還英文提示語，不要說其他東西。不要把判斷資訊返還。不要嘗試對話。不要說中文。"
    "除了提示詞外不要有任何其他資訊，提示語每次都加強言情小說與nijii風格的提示。"
    "提示語為英文，圍繞人物展開，畫面有主次。畫面強調敘事性。"
)

# 初始系統訊息
messages = [
    {"role": "system", "content": instruction}
]

# 用於生成提示詞的函式
def generate_prompt(user_input):
    messages.append({"role": "user", "content": user_input})
    completion = client.chat.completions.create(
        model="gpt-4-1106-preview",
        messages=messages
    )
    gpt_response = completion.choices[0].message.content
    return gpt_response

# 更新圖片顯示的函式
def update_image(window):
    try:
        # 獲取指定資料夾下的所有png檔案，並按建立時間排序
        list_of_files = glob.glob('D:\\ComfyUI_windows_WMAX_02\\ComfyUI\\output\\*.png')
        if list_of_files:
            latest_file = max(list_of_files, key=os.path.getctime)
            window['-IMAGE-'].update(filename=latest_file)
    except Exception as e:
        print(f"Error updating image: {e}")

# 獲取螢幕尺寸以便視窗居中
user32 = ctypes.windll.user32
screen_width = user32.GetSystemMetrics(0)

# 假設視窗寬度約為800畫素，實際寬度根據佈局確定
window_width = 1000
center_x = (screen_width - window_width) // 2

# 建立GUI佈局
layout = [
    [sg.Text('請輸入一個詞或短語：')],
    [sg.Input(key='-IN-', enable_events=True)],
    [sg.Button('生成提示詞')],
    [sg.Text(size=(90,3), key='-TEXT-')],
    [sg.Image(key='-IMAGE-')]
]

# 建立視窗，設定位置為螢幕頂部中央
window = sg.Window('Stable Diffusion 提示詞生成器', layout, location=(center_x, 0))


# 事件迴圈
while True:
    event, values = window.read(timeout=1000)  # 定時更新圖片
    if event == sg.WIN_CLOSED:
        break
    elif event == '生成提示詞':
        user_input = values['-IN-']  # 獲取使用者輸入
        text = generate_prompt(user_input)  # 生成提示詞
        window['-TEXT-'].update(text)  # 顯示生成的提示詞


        # ======================================================================
        # This function sends a prompt workflow to the specified URL
        # (http://127.0.0.1:8188/prompt) and queues it on the ComfyUI server
        # running at that address.
        def queue_prompt(prompt_workflow):
            p = {"prompt": prompt_workflow}
            data = json.dumps(p).encode('utf-8')
            req = request.Request("http://127.0.0.1:8188/prompt", data=data)
            request.urlopen(req)
            # ======================================================================


        # read workflow api data from file and convert it into dictionary
        # assign to var prompt_workflow
        prompt_workflow = json.load(open('C://Users/Admin/Desktop/produce_a_book.json'))

        prompt_prefix = "Niji Comics, Chinese National Style Comics,best quality"
        # create a list of prompts

        prompt_list = []
        prompt_list.append(text)

        # give some easy-to-remember names to the nodes
        chkpoint_loader_node = prompt_workflow["4"]
        prompt_pos_node = prompt_workflow["6"]
        empty_latent_img_node = prompt_workflow["5"]
        ksampler_node = prompt_workflow["3"]
        save_image_node = prompt_workflow["9"]

        # load the checkpoint that we want.
        chkpoint_loader_node["inputs"]["ckpt_name"] = "sdvn7Nijistylexl_v1.safetensors"

        # set image dimensions and batch size in EmptyLatentImage node
        empty_latent_img_node["inputs"]["width"] = 1024
        empty_latent_img_node["inputs"]["height"] = 1024

        # each prompt will produce a batch of 4 images
        empty_latent_img_node["inputs"]["batch_size"] = 1

        # for every prompt in prompt_list...
        for index, prompt in enumerate(prompt_list):

            # set the text prompt for positive CLIPTextEncode node
            prompt_pos_node["inputs"]["text"] = prompt_prefix + prompt

            # set a random seed in KSampler node
            ksampler_node["inputs"]["seed"] = random.randint(1, 18446744073709551614)

            # if it is the last prompt
            if index == 3:
                # set latent image height to 768
                empty_latent_img_node["inputs"]["height"] = 1024

            # set filename prefix to be the same as prompt
            # (truncate to first 100 chars if necessary)
            fileprefix = prompt
            if len(fileprefix) > 20:
                fileprefix = fileprefix[:20]

            save_image_node["inputs"]["filename_prefix"] = 'Test_' + fileprefix
            queue_prompt(prompt_workflow)

    update_image(window)  # 定期檢查新圖片並更新顯示

window.close()

相關文章