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 327c4e1..2e99a3e 100644 Binary files a/requirements.txt and b/requirements.txt differ diff --git a/workflow.json b/workflow.json deleted file mode 100644 index 3c57cf9..0000000 --- a/workflow.json +++ /dev/null @@ -1,107 +0,0 @@ -{ - "3": { - "inputs": { - "seed": 513041955012511, - "steps": 6, - "cfg": 1.5, - "sampler_name": "dpmpp_2m_sde", - "scheduler": "karras", - "denoise": 1, - "model": [ - "4", - 0 - ], - "positive": [ - "6", - 0 - ], - "negative": [ - "7", - 0 - ], - "latent_image": [ - "5", - 0 - ] - }, - "class_type": "KSampler", - "_meta": { - "title": "KSampler" - } - }, - "4": { - "inputs": { - "ckpt_name": "realisticVisionV60B1_v51HyperVAE.safetensors" - }, - "class_type": "CheckpointLoaderSimple", - "_meta": { - "title": "Load Checkpoint" - } - }, - "5": { - "inputs": { - "width": 512, - "height": 512, - "batch_size": 1 - }, - "class_type": "EmptyLatentImage", - "_meta": { - "title": "Empty Latent Image" - } - }, - "6": { - "inputs": { - "text": "Ultra realistic headshot with transparent background of a male soccer player looking at the camera being sixteen years old from Serbia with no facial hair, and dreadlocks hair", - "clip": [ - "4", - 1 - ] - }, - "class_type": "CLIPTextEncode", - "_meta": { - "title": "CLIP Text Encode (Prompt)" - } - }, - "7": { - "inputs": { - "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", - "clip": [ - "4", - 1 - ] - }, - "class_type": "CLIPTextEncode", - "_meta": { - "title": "CLIP Text Encode (Prompt)" - } - }, - "8": { - "inputs": { - "samples": [ - "3", - 0 - ], - "vae": [ - "4", - 2 - ] - }, - "class_type": "VAEDecode", - "_meta": { - "title": "VAE Decode" - } - }, - "9": { - "inputs": { - "filename_prefix": "ComfyUI", - "images": [ - "8", - 0 - ] - }, - "class_type": "SaveImage", - "_meta": { - "title": "Save Image" - } - } -} \ No newline at end of file diff --git a/workflow_api.json b/workflow_api.json new file mode 100644 index 0000000..078f55e --- /dev/null +++ b/workflow_api.json @@ -0,0 +1,107 @@ +{ + "3": { + "inputs": { + "seed": 531046367, + "steps": 6, + "cfg": 1.5, + "sampler_name": "dpmpp_2m_sde", + "scheduler": "karras", + "denoise": 1, + "model": [ + "4", + 0 + ], + "positive": [ + "6", + 0 + ], + "negative": [ + "7", + 0 + ], + "latent_image": [ + "5", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "4": { + "inputs": { + "ckpt_name": "realisticVisionV60B1_v51HyperVAE.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "5": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "6": { + "inputs": { + "text": "Ultra realistic headshot with transparent background of a Olive skinned male soccer player looking at the camera being thirty-three years old from Brazil with moustache, and Bald Medium Blonde long hair", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "positive" + } + }, + "7": { + "inputs": { + "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", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "8": { + "inputs": { + "samples": [ + "3", + 0 + ], + "vae": [ + "4", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "9": { + "inputs": { + "filename_prefix": "FM", + "images": [ + "8", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + } +} \ No newline at end of file