rework with module

This commit is contained in:
Karl Hudgell 2024-12-13 10:46:31 +00:00
parent a049300243
commit 29e6d85a83
6 changed files with 195 additions and 288 deletions

1
.vscode/launch.json vendored
View File

@ -10,6 +10,7 @@
"request": "launch", "request": "launch",
"program": "${file}", "program": "${file}",
"console": "integratedTerminal", "console": "integratedTerminal",
"justMyCode": false,
"args": [ "args": [
"--num_inference_steps", "--num_inference_steps",
"6", "6",

View File

@ -1,223 +1,129 @@
""" import argparse
"""
from lib.rtf_parser import RTF_Parser
import aiohttp
import asyncio
import json
from PIL import Image
from io import BytesIO
import os import os
import random import random
import json import json
import configparser
import pycountry import pycountry
import inflect import inflect
import configparser from lib.rtf_parser import RTF_Parser
import argparse 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 = configparser.ConfigParser()
user_config.read('config.cfg') user_config.read("config.cfg")
rtf = RTF_Parser() rtf = RTF_Parser()
p = inflect.engine() p = inflect.engine()
async def generate_image(session, prompt, num_inference_steps, model): def generate_image(uid,comfy_prompt,app_config,model):
import random """Generate an image using the Comfy API."""
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]},
},
}
}
)
url = user_config["general"]["url"] 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): # Set workflow parameters
# if image_bytes: wf.set_node_param("KSampler", "seed", random.getrandbits(32))
# os.makedirs(folder, exist_ok=True) wf.set_node_param("KSampler", "steps", 6)
# full_path = os.path.join(folder, filename) wf.set_node_param("positive", "text", comfy_prompt)
wf.set_node_param("Save Image", "filename_prefix", uid)
# # Open the image using PIL wf.set_node_param("Load Checkpoint", "ckpt_name", model)
# image = Image.open(BytesIO(image_bytes))
# queue your workflow for completion
# # Resize the image if the resize option is True results = api.queue_and_wait_images(wf, "Save Image")
# if resize: for filename, image_data in results.items():
# image = image.resize((256, 256), Image.LANCZOS) with open(f"./generated_images/{uid}.png", "wb+") as f:
f.write(image_data)
# # Save the image
# image.save(full_path)
# print(f"Image saved as {full_path}")
# else:
# print("Failed to generate or save the image.")
def get_country_name_from_code(code): def get_country_name_from_code(code):
"""Get country name from 3-letter ISO code."""
try: try:
country = pycountry.countries.get(alpha_3=code.upper()) country = pycountry.countries.get(alpha_3=code.upper())
return country.name if country else "Unknown country code" return country.name if country else "Unknown country code"
except KeyError: except KeyError:
return "Invalid country code" 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 = [] def generate_prompts_for_players(players, app_config):
for i in range(1): """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]) country = get_country_name_from_code(player[1])
facial_characteristics = random.choice(config["facial_characteristics"]) facial_characteristics = random.choice(app_config["facial_characteristics"])
hair_length = config["hair_length"][player[5]] hair_length = app_config["hair_length"][player[5]]
hair_colour = config["hair_color"][player[6]] hair_colour = app_config["hair_color"][player[6]]
skin_tone = config["skin_tone_map"][player[7]] skin_tone = app_config["skin_tone_map"][player[7]]
player_age = p.number_to_words(player[3]) 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, skin_tone=skin_tone,
age=player_age, age=player_age,
country=country, country=country,
facial_characteristics=( facial_characteristics=facial_characteristics or "no facial hair",
facial_characteristics if facial_characteristics else "no facial hair"
),
hair=f"{hair_length} {hair_colour} {hair}", hair=f"{hair_length} {hair_colour} {hair}",
) )
print(f"Generated prompt: {prompt}") print(f"Generated prompt: {prompt}")
prompt = f"{player[0]}:{prompt}"
prompts.append(prompt)
return prompts
task = asyncio.create_task( def main():
generate_image( """Main function for generating images."""
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():
parser = argparse.ArgumentParser(description="Generate images for country groups") 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(
parser.add_argument("--num_inference_steps", type=int, default=6, help="Number of inference steps. Defaults to 1") "--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() args = parser.parse_args()
if not args.rtf_file: if not args.rtf_file:
raise Exception("Please pass in a RTF file as --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: with open("config.json", "r") as f:
app_config = json.load(f) app_config = json.load(f)
total_images = len(rtf_file) prompts = generate_prompts_for_players(rtf_file, app_config)
# if not ask_user_confirmation(total_images, total_cost): for prompt in prompts:
# print("Operation cancelled by user.") uid = prompt.split(":")[0]
# sys.exit(0) comfy_prompt = prompt.split(":")[1]
generate_image(
generated_images = 0 uid,
async with aiohttp.ClientSession() as session: comfy_prompt,
tasks = [] app_config,
for player in rtf_file[:500]: user_config["general"]["model"],
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)
print("\nImage generation complete for all country groups.") print("\nImage generation complete for all country groups.")
print(f"Total images generated: {generated_images}")
if __name__ == "__main__": if __name__ == "__main__":
asyncio.run(main()) main()

View File

@ -1,7 +1,7 @@
# Football Manager AI Face Generator # ComfyUI FM NewGen
## Overview
<!-- ## Overview -->
<!--
This tool generates AI faces for Football Manager, compatible with NewGan and similar AI face pack installers (FMRTE). It uses the DeepInfra API to create realistic player faces based on various ethnic groups and facial characteristics. The image generation model is Flux ('schnell' by default). You can obtain your DeepInfra API key from [here](https://deepinfra.com/). This tool generates AI faces for Football Manager, compatible with NewGan and similar AI face pack installers (FMRTE). It uses the DeepInfra API to create realistic player faces based on various ethnic groups and facial characteristics. The image generation model is Flux ('schnell' by default). You can obtain your DeepInfra API key from [here](https://deepinfra.com/).
Short tutorial on how to use: https://www.youtube.com/watch?v=vldujOEPWoo Short tutorial on how to use: https://www.youtube.com/watch?v=vldujOEPWoo
@ -117,4 +117,4 @@ For issues or questions, please open an issue on the GitHub repository.
## License ## License
This project is open-source and available under the MIT License. This project is open-source and available under the MIT License. -->

Binary file not shown.

View File

@ -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"
}
}
}

107
workflow_api.json Normal file
View File

@ -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"
}
}
}