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",
"program": "${file}",
"console": "integratedTerminal",
"justMyCode": false,
"args": [
"--num_inference_steps",
"6",

View File

@ -1,114 +1,49 @@
"""
"""
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"
@ -116,108 +51,79 @@ def get_country_name_from_code(code):
return "Invalid country code"
async def generate_images_for_country_group(
session,
player,
config,
model,
):
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)
tasks = []
for i in range(1):
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")
# Parse the 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])
)
# 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,
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"],
)
)
tasks.append(task)
results = await asyncio.gather(*tasks)
# generated_images = sum(results)
print("\nImage generation complete for all country groups.")
print(f"Total images generated: {generated_images}")
if __name__ == "__main__":
asyncio.run(main())
main()

View File

@ -1,7 +1,7 @@
# Football Manager AI Face Generator
## Overview
# ComfyUI FM NewGen
<!-- ## 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/).
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
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"
}
}
}