From 29e6d85a83a7d52bfd93279dd902585e345bb8c8 Mon Sep 17 00:00:00 2001 From: Karl Hudgell Date: Fri, 13 Dec 2024 10:46:31 +0000 Subject: [PATCH] rework with module --- .vscode/launch.json | 1 + comfy_fm_newgen.py | 260 ++++++++++++++------------------------------ readme.md | 8 +- requirements.txt | Bin 43 -> 1788 bytes workflow.json | 107 ------------------ workflow_api.json | 107 ++++++++++++++++++ 6 files changed, 195 insertions(+), 288 deletions(-) delete mode 100644 workflow.json create mode 100644 workflow_api.json diff --git a/.vscode/launch.json b/.vscode/launch.json index fe0f916..e577a69 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -10,6 +10,7 @@ "request": "launch", "program": "${file}", "console": "integratedTerminal", + "justMyCode": false, "args": [ "--num_inference_steps", "6", diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index 50f632f..00d1015 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -1,223 +1,129 @@ -""" -""" -from lib.rtf_parser import RTF_Parser -import aiohttp -import asyncio -import json -from PIL import Image -from io import BytesIO +import argparse import os import random import json +import configparser import pycountry import inflect -import configparser -import argparse +from lib.rtf_parser import RTF_Parser +import logging +import sys +from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper +# logging.basicConfig(stream=sys.stdout, level=logging.INFO) + + +# Load user configurations user_config = configparser.ConfigParser() -user_config.read('config.cfg') +user_config.read("config.cfg") rtf = RTF_Parser() p = inflect.engine() -async def generate_image(session, prompt, num_inference_steps, model): - import random - - payload = json.dumps( - { - "prompt": { - "3": { - "class_type": "KSampler", - "inputs": { - "cfg": 1.5, - "denoise": 1, - "latent_image": ["5", 0], - "model": ["4", 0], - "negative": ["7", 0], - "positive": ["6", 0], - "sampler_name": "dpmpp_2m_sde", - "scheduler": "karras", - "seed": random.getrandbits(32), - "steps": num_inference_steps, - }, - }, - "4": { - "class_type": "CheckpointLoaderSimple", - "inputs": {"ckpt_name": model}, - }, - "5": { - "class_type": "EmptyLatentImage", - "inputs": {"batch_size": 1, "height": 512, "width": 512}, - }, - "6": { - "class_type": "CLIPTextEncode", - "inputs": { - "clip": ["4", 1], - "text": prompt, - }, - }, - "7": { - "class_type": "CLIPTextEncode", - "inputs": { - "clip": ["4", 1], - "text": "(nsfw, naked, nude, deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation", - }, - }, - "8": { - "class_type": "VAEDecode", - "inputs": {"samples": ["3", 0], "vae": ["4", 2]}, - }, - "9": { - "class_type": "SaveImage", - "inputs": {"filename_prefix": "FM", "images": ["8", 0]}, - }, - } - } - ) +def generate_image(uid,comfy_prompt,app_config,model): + """Generate an image using the Comfy API.""" url = user_config["general"]["url"] - url = f"{url}prompt" - headers = {"Content-Type": "application/json"} - async with session.post(url, headers=headers, data=payload) as response: - if response.status == 200: - # response_data = await response.json() - # image_data = response_data["images"][0].split(",")[1] - # image_bytes = base64.b64decode(image_data) - # return image_bytes - pass - else: - return None + # Initialize API and workflow + api = ComfyApiWrapper(url) + wf = ComfyWorkflowWrapper("./workflow_api.json") -# def save_image(image_bytes, folder, filename, resize=False): -# if image_bytes: -# os.makedirs(folder, exist_ok=True) -# full_path = os.path.join(folder, filename) - -# # Open the image using PIL -# image = Image.open(BytesIO(image_bytes)) - -# # Resize the image if the resize option is True -# if resize: -# image = image.resize((256, 256), Image.LANCZOS) - -# # Save the image -# image.save(full_path) -# print(f"Image saved as {full_path}") -# else: -# print("Failed to generate or save the image.") - + # Set workflow parameters + wf.set_node_param("KSampler", "seed", random.getrandbits(32)) + wf.set_node_param("KSampler", "steps", 6) + wf.set_node_param("positive", "text", comfy_prompt) + wf.set_node_param("Save Image", "filename_prefix", uid) + wf.set_node_param("Load Checkpoint", "ckpt_name", model) + + # queue your workflow for completion + results = api.queue_and_wait_images(wf, "Save Image") + for filename, image_data in results.items(): + with open(f"./generated_images/{uid}.png", "wb+") as f: + f.write(image_data) def get_country_name_from_code(code): + """Get country name from 3-letter ISO code.""" try: country = pycountry.countries.get(alpha_3=code.upper()) return country.name if country else "Unknown country code" except KeyError: return "Invalid country code" - -async def generate_images_for_country_group( - session, - player, - config, - model, -): - print(f"\nGenerating image for {player[0]} - {player[8]}") - folder_name = f"generated_images/" - os.makedirs(folder_name, exist_ok=True) - tasks = [] - for i in range(1): +def generate_prompts_for_players(players, app_config): + """Generate images for a specific player and configuration.""" + prompts = [] + for player in players[:100]: + print(f"\nGenerating image for {player[0]} - {player[8]}") + folder_name = f"generated_images/" + os.makedirs(folder_name, exist_ok=True) + country = get_country_name_from_code(player[1]) - facial_characteristics = random.choice(config["facial_characteristics"]) - hair_length = config["hair_length"][player[5]] - hair_colour = config["hair_color"][player[6]] - skin_tone = config["skin_tone_map"][player[7]] + facial_characteristics = random.choice(app_config["facial_characteristics"]) + hair_length = app_config["hair_length"][player[5]] + hair_colour = app_config["hair_color"][player[6]] + skin_tone = app_config["skin_tone_map"][player[7]] player_age = p.number_to_words(player[3]) - hair = random.choice(config["hair"]) + hair = random.choice(app_config["hair"]) - prompt = config["prompt"].format( + # Format the prompt + prompt = app_config["prompt"].format( skin_tone=skin_tone, age=player_age, country=country, - facial_characteristics=( - facial_characteristics if facial_characteristics else "no facial hair" - ), + facial_characteristics=facial_characteristics or "no facial hair", hair=f"{hair_length} {hair_colour} {hair}", ) - print(f"Generated prompt: {prompt}") + prompt = f"{player[0]}:{prompt}" + prompts.append(prompt) + return prompts - task = asyncio.create_task( - generate_image( - session=session, - prompt=prompt, - num_inference_steps=6, - model=model, - ) - ) - tasks.append(task) - - image_bytes_list = await asyncio.gather(*tasks) - - # for i, image_bytes in enumerate(image_bytes_list): - # if image_bytes: - # next_number = get_next_image_number(folder_name) - # file_name = f"{country_group}{next_number}.png" - # save_image(image_bytes, folder_name, file_name, resize) - - # return len([img for img in image_bytes_list if img is not None]) - - -async def main(): +def main(): + """Main function for generating images.""" parser = argparse.ArgumentParser(description="Generate images for country groups") - parser.add_argument("--rtf_file", type= str, default=None, help= "Path to the rtf file to be processed") - parser.add_argument("--num_inference_steps", type=int, default=6, help="Number of inference steps. Defaults to 1") - + parser.add_argument( + "--rtf_file", + type=str, + default=None, + help="Path to the RTF file to be processed", + ) + parser.add_argument( + "--num_inference_steps", + type=int, + default=6, + help="Number of inference steps. Defaults to 6", + ) args = parser.parse_args() + if not args.rtf_file: raise Exception("Please pass in a RTF file as --rtf_file") - rtf_file = rtf.parse_rtf(args.rtf_file) - - # Extract unique values from positions 5, 6, and 7 - hair_length = list( - set(item[pos] for item in rtf_file for pos in [5]) - ) - hair_colour = list( - set(item[pos] for item in rtf_file for pos in [6]) - ) - skin_tone = list( - set(item[pos] for item in rtf_file for pos in [7]) - ) + # Parse the RTF file + rtf_file = rtf.parse_rtf(args.rtf_file) + + # Extract unique values + hair_length = list(set(item[5] for item in rtf_file)) + hair_colour = list(set(item[6] for item in rtf_file)) + skin_tone = list(set(item[7] for item in rtf_file)) + + # Load configurations with open("config.json", "r") as f: app_config = json.load(f) - total_images = len(rtf_file) - # if not ask_user_confirmation(total_images, total_cost): - # print("Operation cancelled by user.") - # sys.exit(0) - - generated_images = 0 - async with aiohttp.ClientSession() as session: - tasks = [] - for player in rtf_file[:500]: - task = asyncio.create_task( - generate_images_for_country_group( - session, - player, - app_config, - user_config["general"]["model"], - ) - ) - tasks.append(task) - - results = await asyncio.gather(*tasks) - # generated_images = sum(results) + prompts = generate_prompts_for_players(rtf_file, app_config) + for prompt in prompts: + uid = prompt.split(":")[0] + comfy_prompt = prompt.split(":")[1] + generate_image( + uid, + comfy_prompt, + app_config, + user_config["general"]["model"], + ) print("\nImage generation complete for all country groups.") - print(f"Total images generated: {generated_images}") - if __name__ == "__main__": - asyncio.run(main()) + main() diff --git a/readme.md b/readme.md index 4c171c8..9097f84 100644 --- a/readme.md +++ b/readme.md @@ -1,7 +1,7 @@ -# Football Manager AI Face Generator - -## Overview +# ComfyUI FM NewGen + + diff --git a/requirements.txt b/requirements.txt index 327c4e1d79d98f9f0a0fd3d99df4ea86657b1973..2e99a3ee9fa12135cc2e07a05d6cac49c844fe52 100644 GIT binary patch literal 1788 zcmZvd+m6~m5JcxWQa;5ImxPdqd_tuB0b&8O4)_vm_V)2f>U1&8K17yTOjlP|SNC}T z{@JD`_31Hf)1d1$9pZk_^I6xJzUx|~Pw8{qAJYd#oH30&2{lPO9e9uc}!h! zfdgTzF{5gAX_qQ7HK|qYK3(-4>=$8Y!DuYiE_XesL6xiRlz9n`T0B-fQykP<_wtUT zPgo8@SE}9XsSSNidJC4^54oAWmlGhjfdzpw5Uc2Rr)ZmMb+Zv>kcN$F(9xjGMv+F` z`%s1lD4$1WE!I9NcYz#l-Gj0Nbl)f!7qhnPG8MLR_#i}o>zi&@cm76swU`fb*}i^N ztn~IvI(uP8#UF)&k18M*F)MrUKU}4(A(-9aPPxp{e6TF!lc#{<-1Awux9|#Ax9W}Q z6<6t)-FuQ6+?y0HQTE50T4i4pt8N%Jpx2eR)-{WZt#<1i3iry$e%Pzb*@=%CRMM5V zGkeQSK)X7~DLsdVQ+f%7XtQ(l(5YUVI7wOkDqP688aPqkxfl^3&&c|unelaODBCE)o-UyS1JnGRGjYB2Qi;vIINNFWu|p>*{W(I z$4gE6^$8e+=wpY><8$Y3hS>2zQS`87VBBrq33}udU8R$Ao?Ss}+2)`6{YNNb&vHl1 z#rN{o_j|wgVr7Fmd6~6Q+vkt{sde)gzY8mvWt_o+WSZGU&x?uTHTl1#-rG~&noCU- z9%BYMINl~&JSYfGVWWw$^K5lK-35yuu4ruDANTJ)pcXrP5d+>czsz9Q*P#aw>$ung zPoq!-Y?FH1OARcZ2yYSlGHE8f1&->TcdOrWpo$t}?CnO_7ct^D@5u7DFdjF& sH&?N-osC>yMs3HkpU>iW5C?mA()UTHq)$ literal 43 ycmYey%+Dw(Dd6P_$jr&fFX!bds4U6I&(lrGFG