import json from websocket import create_connection, WebSocketTimeoutException, WebSocketConnectionClosedException import uuid import urllib.request import urllib.parse import random import time from datetime import datetime # 定义一个函数来显示GIF图片 def show_gif(fname): import base64 from IPython import display with open(fname, 'rb') as fd: b64 = base64.b64encode(fd.read()).decode('ascii') return display.HTML(f'') # 定义一个函数向服务器队列发送提示信息 def queue_prompt(prompt): p = {"prompt": prompt, "client_id": client_id} data = json.dumps(p).encode('utf-8') req = urllib.request.Request("http://{}/prompt".format(server_address), data=data) return json.loads(urllib.request.urlopen(req).read()) # 定义一个函数来获取图片 def get_image(filename, subfolder, folder_type): data = {"filename": filename, "subfolder": subfolder, "type": folder_type} url_values = urllib.parse.urlencode(data) with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response: return response.read() # 定义一个函数来获取历史记录 def get_history(prompt_id): with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response: return json.loads(response.read()) # 定义一个函数来获取图片，这涉及到监听WebSocket消息 def get_images(ws, prompt): prompt_id = queue_prompt(prompt)['prompt_id'] print('Prompt: ', prompt) print('Prompt ID: ', prompt_id) output_images = {} # 等待执行完成 start_time = time.time() while True: if time.time() - start_time > 1200: # 设置2分钟超时 print("超时：等待执行完成超过120秒") break out = ws.recv() if isinstance(out, str): message = json.loads(out) if message['type'] == 'executing': data = message['data'] if data['node'] is None and data['prompt_id'] == prompt_id: print('Execution complete') break # 执行完成 else: continue # 预览为二进制数据 # 获取完成的历史记录 history = get_history(prompt_id).get(prompt_id, {}) if not history: print("未找到该提示的历史记录") return {} for node_id in history['outputs']: node_output = history['outputs'][node_id] # 图片分支 if 'images' in node_output: images_output = [] for image in node_output['images']: try: image_data = get_image(image['filename'], image['subfolder'], image['type']) images_output.append(image_data) except Exception as e: print(f"获取图像错误: {str(e)}") output_images[node_id] = images_output # 视频分支 elif 'videos' in node_output: videos_output = [] for video in node_output['videos']: try: video_data = get_image(video['filename'], video['subfolder'], video['type']) videos_output.append(video_data) except Exception as e: print(f"获取视频错误: {str(e)}") output_images[node_id] = videos_output print(f'Obtained {len(output_images)} image/video sets') return output_images # 解析工作流并获取图片 def parse_workflow(prompt, seed, workflowfile): print(f'Workflow file: {workflowfile}') try: with open(workflowfile, 'r', encoding="utf-8") as f: prompt_data = json.load(f) # 设置文本提示 prompt_data["6"]["inputs"]["text"] = prompt # 设置随机种子（如果需要） if "Ksampler" in prompt_data: if "seed" in prompt_data["Ksampler"]["inputs"]: prompt_data["Ksampler"]["inputs"]["seed"] = seed elif "noise_seed" in prompt_data["Ksampler"]["inputs"]: prompt_data["Ksampler"]["inputs"]["noise_seed"] = seed return prompt_data except Exception as e: print(f"工作流解析错误: {str(e)}") return {} # 生成图像并保存 def generate_clip(prompt, seed, workflowfile, idx): print(f'Processing prompt #{idx}: "{prompt[:50]}{"..." if len(prompt) > 50 else ""}"') print(f'Using seed: {seed}') try: # 使用正确的WebSocket连接方式 ws_url = f"ws://{server_address}/ws?clientId={client_id}" # 使用 create_connection ws = create_connection(ws_url, timeout=600) # 解析工作流 workflow_data = parse_workflow(prompt, seed, workflowfile) if not workflow_data: print("工作流数据为空") return # 获取图像 images = get_images(ws, workflow_data) # 关闭连接 ws.close() except WebSocketTimeoutException as e: print(f"WebSocket连接超时: {str(e)}") return except WebSocketConnectionClosedException as e: print(f"WebSocket连接已关闭: {str(e)}") return except Exception as e: print(f"WebSocket错误: {str(e)}") return saved_files = [] if images: for node_id, image_list in images.items(): for i, image_data in enumerate(image_list): # 格式化时间戳 timestamp = datetime.now().strftime("%Y%m%d%H%M%S") # 创建唯一文件名 filename = f"{idx}_{seed}_{timestamp}_{i}.png" file_path = f"{WORKING_DIR}/{filename}" print(f'Saving to: {file_path}') try: with open(file_path, "wb") as f: f.write(image_data) saved_files.append(file_path) except Exception as e: print(f"文件保存错误: {str(e)}") else: print("未获取到图像数据") if saved_files: print(f"成功生成 {len(saved_files)} 张图片") else: print("未保存任何图片") # 直接在代码中定义提示词列表 PROMPTS = [ "A beautiful sunset over the ocean, realistic, cinematic lighting", "A futuristic cityscape at night, cyberpunk style, neon lights", "An ancient forest with magical creatures, fantasy, photorealistic", "A steampunk laboratory with bubbling beakers and intricate machinery", "Abstract geometric patterns in vibrant colors, digital art", "A majestic lion in the African savannah, golden hour lighting", "A cozy cabin in the mountains during winter, warm lights inside", "A detailed close-up of a butterfly on a flower, macro photography", "Underwater scene with coral reef and tropical fish, crystal clear water" ] if __name__ == "__main__": # 设置工作目录和项目相关的路径 WORKING_DIR = 'output' workflowfile = './workflows/flux_redux.json' COMFYUI_ENDPOINT = '127.0.0.1:8188' # 服务器配置 global server_address, client_id server_address = COMFYUI_ENDPOINT client_id = str(uuid.uuid4()) # 生成一个唯一的客户端ID # 种子设置 - 可以选择固定或随机 USE_RANDOM_SEED = True # 设为False则使用固定种子 base_seed = 15465856 print(f"Starting image generation with ComfyUI at {server_address}") print(f"Working directory: {WORKING_DIR}") print(f"Workflow file: {workflowfile}") print(f"Client ID: {client_id}") start_time = time.time() # 处理每个提示词 for idx, prompt in enumerate(PROMPTS, start=1): # 设置种子 current_seed = random.randint(1, 10**8) if USE_RANDOM_SEED else base_seed print(f"\n===== Processing Prompt #{idx} of {len(PROMPTS)} =====") generate_clip(prompt, current_seed, workflowfile, idx) # 添加延迟以避免服务器过载 time.sleep(2) # 2秒延迟 elapsed = time.time() - start_time print(f"\nProcessing completed in {elapsed:.2f} seconds") print(f"Generated images for {len(PROMPTS)} prompts")