rework commit

This commit is contained in:
Karl Hudgell 2024-12-13 09:35:47 +00:00
parent 43054a0a12
commit 9b76e8bf1a
25 changed files with 16616 additions and 241 deletions

1
.env
View File

@ -1 +0,0 @@
DEEPINFRA_API_KEY = "YOUR API KEY HERE"

3
.gitignore vendored
View File

@ -1 +1,2 @@
venv/*
venv/*
config.cfg

10
.vscode/launch.json vendored
View File

@ -11,14 +11,10 @@
"program": "${file}",
"console": "integratedTerminal",
"args": [
"--n_per_country",
"2",
"--width",
"1024",
"--height",
"1024",
"--num_inference_steps",
"2"
"6",
"--rtf_file",
"./2030New.rtf"
]
}
]

16049
2030New.rtf Normal file

File diff suppressed because it is too large Load Diff

Binary file not shown.

View File

@ -1,190 +0,0 @@
"""
Author: Emil Mirzayev
This script generates images based on prompts for different country groups using the DeepInfra API.
It allows for customization of image parameters and provides cost estimation before execution.
Features:
- Asynchronous image generation for improved performance
- Configurable number of images per country group
- Customizable image dimensions and inference steps
- Option to resize generated images
- Cost estimation and user confirmation before execution
- Saves generated images in separate folders for each country group
Usage:
python async_image_generator.py [--n_per_country N] [--width W] [--height H]
[--num_inference_steps S] [--model MODEL] [--resize]
Requirements:
- Python 3.7+
- Required packages: aiohttp, Pillow, python-dotenv
- DeepInfra API key (set in .env file)
- config.json file with country groups, prompts, and other configuration details
Note: Ensure you have sufficient API credits before running large batches.
"""
import aiohttp
import asyncio
import json
import base64
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv
import os
import random
import argparse
import sys
# Load the .env file from the current directory
load_dotenv()
DEEPINFRA_API_KEY = os.getenv('DEEPINFRA_API_KEY')
async def generate_image(session, api_key, prompt, width, height, num_inference_steps, model):
if model == "schnell":
url = "https://api.deepinfra.com/v1/inference/black-forest-labs/FLUX-1-schnell"
elif model == "dev":
url ="https://api.deepinfra.com/v1/inference/black-forest-labs/FLUX-1-dev"
else:
return None
headers = {'Authorization': f'Bearer {api_key}'}
data = {"prompt": prompt, "width": width, "height": height, "num_inference_steps": num_inference_steps}
json_data = json.dumps(data)
async with session.post(url, headers=headers, data=json_data) 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
else:
return None
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.")
def get_next_image_number(folder):
os.makedirs(folder, exist_ok=True)
existing_files = [f for f in os.listdir(folder) if f.endswith('.png')]
return len(existing_files) + 1
def calculate_total_images(config, n_per_country):
return len(config['countries']) * n_per_country
def ask_user_confirmation(total_images, total_cost):
print(f"\nPotential cost of this run: ${total_cost:.4f} for {total_images} images")
while True:
response = input("Do you want to proceed? (yes/no): ").lower().strip()
if response in ['yes', 'y']:
return True
elif response in ['no', 'n']:
return False
else:
print("Please answer with 'yes' or 'no'.")
def calculate_cost(width, height, num_inference_steps):
return 0.0005 * (width / 1024) * (height / 1024) * num_inference_steps
async def generate_images_for_country_group(session, country_group, config, n_per_country, width, height, num_inference_steps, model, resize):
print(f"\nGenerating images for {country_group}")
folder_name = f"generated_images/{country_group}"
os.makedirs(folder_name, exist_ok=True)
tasks = []
for i in range(n_per_country):
country = random.choice(config['countries'][country_group])
facial_characteristics = random.choice(config['facial_characteristics'])
hair = random.choice(config['hair'])
prompt = config['prompt'].format(
country=country,
facial_characteristics=facial_characteristics if facial_characteristics else "no facial hair",
hair=hair
)
print(f"Generated prompt: {prompt}")
task = asyncio.create_task(generate_image(
session=session,
api_key=DEEPINFRA_API_KEY,
prompt=prompt,
width=width,
height=height,
num_inference_steps=num_inference_steps,
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.add_argument("--n_per_country", type=int, default=1, help="Number of images to generate per country group. Defaults to 1")
parser.add_argument("--width", type=int, default=512, help="Width of the generated images. Defaults to 512")
parser.add_argument("--height", type=int, default=512, help="Height of the generated images. Defaults to 512")
parser.add_argument("--num_inference_steps", type=int, default=1, help="Number of inference steps. Defaults to 1")
parser.add_argument("--model", type= str, default= "schnell", help= "The model to be used. Must be one of `schnell` or `dev`. Schnell is cheaper and faster")
parser.add_argument("--resize", action="store_true", help="Resize images to 256x256 if set")
args = parser.parse_args()
with open('config.json', 'r') as f:
config = json.load(f)
total_images = calculate_total_images(config, args.n_per_country)
total_cost = total_images * calculate_cost(args.width, args.height, args.num_inference_steps)
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 country_group in config['countries'].keys():
task = asyncio.create_task(generate_images_for_country_group(
session, country_group, config, args.n_per_country,
args.width, args.height, args.num_inference_steps,
args.model, args.resize
))
tasks.append(task)
results = await asyncio.gather(*tasks)
generated_images = sum(results)
actual_cost = generated_images * calculate_cost(args.width, args.height, args.num_inference_steps)
print("\nImage generation complete for all country groups.")
print(f"Total images generated: {generated_images}")
print(f"Actual cost of this run: ${actual_cost:.4f}")
if __name__ == "__main__":
asyncio.run(main())

223
comfy_fm_newgen.py Normal file
View File

@ -0,0 +1,223 @@
"""
"""
from lib.rtf_parser import RTF_Parser
import aiohttp
import asyncio
import json
from PIL import Image
from io import BytesIO
import os
import random
import json
import pycountry
import inflect
import configparser
import argparse
user_config = configparser.ConfigParser()
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]},
},
}
}
)
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
# 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.")
def get_country_name_from_code(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):
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]]
player_age = p.number_to_words(player[3])
hair = random.choice(config["hair"])
prompt = config["prompt"].format(
skin_tone=skin_tone,
age=player_age,
country=country,
facial_characteristics=(
facial_characteristics if facial_characteristics else "no facial hair"
),
hair=f"{hair_length} {hair_colour} {hair}",
)
print(f"Generated prompt: {prompt}")
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():
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")
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])
)
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)
print("\nImage generation complete for all country groups.")
print(f"Total images generated: {generated_images}")
if __name__ == "__main__":
asyncio.run(main())

3
config.cfg.sample Normal file
View File

@ -0,0 +1,3 @@
[general]
url =
model = realisticVisionV60B1_v51HyperVAE.safetensors

View File

@ -1,41 +1,195 @@
{"countries":{
"African": ["Nigeria", "Kenya", "South Africa", "Egypt", "Ghana"],
"Asian": ["China", "Japan", "India", "South Korea", "Vietnam"],
"Caucasian": ["United States", "Canada", "Australia", "New Zealand", "United Kingdom"],
"Central European": ["Germany", "Poland", "Czech Republic", "Austria", "Hungary"],
"EECA": ["Ukraine", "Kazakhstan", "Belarus", "Georgia", "Azerbaijan"],
"Italmed": ["Italy", "Greece", "Croatia", "Malta", "Cyprus"],
"MENA": ["Saudi Arabia", "UAE", "Iran", "Israel", "Morocco"],
"MESA": ["Pakistan", "Afghanistan", "Bangladesh", "Nepal", "Sri Lanka"],
"SAMed": ["Tunisia", "Algeria", "Libya", "Lebanon", "Jordan"],
"Scandinavian": ["Sweden", "Norway", "Denmark", "Finland", "Iceland"],
"Seasian": ["Indonesia", "Malaysia", "Thailand", "Philippines", "Singapore"],
"South American": ["Brazil", "Argentina", "Colombia", "Peru", "Chile"],
"SpanMed": ["Spain", "Portugal", "Andorra", "Gibraltar", "Monaco"],
"YugoGreek": ["Serbia", "Greece", "North Macedonia", "Montenegro", "Albania"]
},
"facial_characteristics": [
"beard",
"moustache",
"clean-shaven",
"stubble",
"goatee",
"sideburns",
null
],
"hair": [
"short",
"long",
"bald",
"buzz cut",
"medium-length",
"curly",
"wavy",
"spiky",
"afro",
"dreadlocks",
"mohawk",
"ponytail"
],
"prompt": "Ultra realistic headshot with transparent background of a male soccer player looking at the camera being twenty five years old from {country} with {facial_characteristics}, and {hair} hair"
}
{
"countries": {
"African": [
"Nigeria",
"Kenya",
"South Africa",
"Egypt",
"Ghana"
],
"Asian": [
"China",
"Japan",
"India",
"South Korea",
"Vietnam"
],
"Caucasian": [
"United States",
"Canada",
"Australia",
"New Zealand",
"United Kingdom"
],
"Central European": [
"Germany",
"Poland",
"Czech Republic",
"Austria",
"Hungary"
],
"EECA": [
"Ukraine",
"Kazakhstan",
"Belarus",
"Georgia",
"Azerbaijan"
],
"Italmed": [
"Italy",
"Greece",
"Croatia",
"Malta",
"Cyprus"
],
"MENA": [
"Saudi Arabia",
"UAE",
"Iran",
"Israel",
"Morocco"
],
"MESA": [
"Pakistan",
"Afghanistan",
"Bangladesh",
"Nepal",
"Sri Lanka"
],
"SAMed": [
"Tunisia",
"Algeria",
"Libya",
"Lebanon",
"Jordan"
],
"Scandinavian": [
"Sweden",
"Norway",
"Denmark",
"Finland",
"Iceland"
],
"Seasian": [
"Indonesia",
"Malaysia",
"Thailand",
"Philippines",
"Singapore"
],
"South American": [
"Brazil",
"Argentina",
"Colombia",
"Peru",
"Chile"
],
"SpanMed": [
"Spain",
"Portugal",
"Andorra",
"Gibraltar",
"Monaco"
],
"YugoGreek": [
"Serbia",
"Greece",
"North Macedonia",
"Montenegro",
"Albania"
]
},
"facial_characteristics": [
"beard",
"moustache",
"clean-shaven",
"stubble",
"goatee",
"sideburns",
null
],
"hair": [
"short",
"long",
"bald",
"buzz cut",
"medium-length",
"curly",
"wavy",
"spiky",
"afro",
"dreadlocks",
"mohawk",
"ponytail"
],
"skin_tone_map": {
"0": "Lightest",
"1": "Lightest",
"2": "Very Light",
"3": "Light",
"4": "Light with Beige Undertones",
"5": "Beige",
"6": "Light Tan",
"7": "Tan",
"8": "Medium Tan",
"9": "Olive",
"10": "Light Brown",
"11": "Brown",
"12": "Medium Brown",
"13": "Dark Brown",
"14": "Very Dark Brown",
"15": "Deep Brown",
"16": "Light Ebony",
"17": "Ebony",
"18": "Dark Ebony",
"19": "Deep Ebony",
"20": "Darkest"
},
"hair_color": {
"0": "Jet Black",
"1": "Jet Black",
"2": "Dark Black",
"3": "Very Dark Brown",
"4": "Dark Brown",
"5": "Medium Brown",
"6": "Light Brown",
"7": "Ash Brown",
"8": "Brown with Red Tint",
"9": "Auburn",
"10": "Chestnut",
"11": "Dark Blonde",
"12": "Medium Blonde",
"13": "Ash Blonde",
"14": "Light Blonde",
"15": "Platinum Blonde",
"16": "Strawberry Blonde",
"17": "Light Ginger",
"18": "Ginger",
"19": "Bright Red",
"20": "Copper Red"
},
"hair_length": {
"0": "Bald",
"1": "Bald",
"2": "Very Short (buzz cut)",
"3": "Short (crew cut)",
"4": "Short with fade",
"5": "Classic Short",
"6": "Medium-Short",
"7": "Medium (textured)",
"8": "Medium (layered)",
"9": "Medium with fringe",
"10": "Medium-Long (shaggy)",
"11": "Long (shoulder-length)",
"12": "Long with parting",
"13": "Long (tied back)",
"14": "Very Long (untied)",
"15": "Curly, Short",
"16": "Curly, Medium",
"17": "Curly, Long",
"18": "Afro-Short",
"19": "Afro-Medium",
"20": "Afro-Long"
},
"prompt": "Ultra realistic headshot with transparent background of a {skin_tone} skinned male soccer player looking at the camera being {age} years old from {country} with {facial_characteristics}, and {hair} hair"
}

Binary file not shown.

Before

Width:  |  Height:  |  Size: 287 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 270 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 264 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 259 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 282 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 281 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 278 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 323 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 262 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 246 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 292 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 242 KiB

View File

@ -1 +0,0 @@
Here are some examples of the script runs

Binary file not shown.

34
lib/rtf_parser.py Normal file
View File

@ -0,0 +1,34 @@
import re
class RTF_Parser:
def __init__(self):
pass
def parse_rtf(self, path):
UID_regex = re.compile('([0-9]){4,}')
result_data = []
rtf = open(path, 'r', encoding="UTF-8")
# self.logger.info(rtf)
rtf_data = []
for line in rtf:
if UID_regex.search(line):
# self.logger.info(line.strip())
rtf_data.append(line.strip())
for newgen in rtf_data:
data_fields = newgen.split('|')
sec_nat = data_fields[3].strip()
if sec_nat == '':
sec_nat = None
result_data.append([data_fields[1].strip(), data_fields[2].strip(), sec_nat, data_fields[5].strip(),data_fields[6].strip(),data_fields[7].strip(),data_fields[8].strip(),data_fields[9].strip(), data_fields[4].strip()])
rtf.close()
return result_data
def is_rtf_valid(self, path):
rtf_regex = re.compile('(\|\s*[0-9]{8,}\s*)(\|\s*([A-Z]{3})*\s*)+(\|[\s*\w*\.*\-*]+)(\|[\s*\d+]+){3}\|')
rtf = open(path, 'r', encoding="UTF-8")
rtf_data = rtf.read()
rtf.close()
if rtf_regex.search(rtf_data):
return True
return False

107
workflow.json Normal file
View File

@ -0,0 +1,107 @@
{
"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"
}
}
}