From 474a96b3f850fea8eea9c9a045dfeb7864de8e4d Mon Sep 17 00:00:00 2001
From: Emil Mirzayev <25964049+emilmirzayev@users.noreply.github.com>
Date: Fri, 11 Oct 2024 15:35:48 +0100
Subject: [PATCH] initial commit

first commit
---
 async_image_generator.py | 190 +++++++++++++++++++++++++++++++++++++++
 config.json              |  41 +++++++++
 readme.md                | 106 ++++++++++++++++++++++
 remove_bg.py             |  75 ++++++++++++++++
 requirements.txt         |   5 ++
 5 files changed, 417 insertions(+)
 create mode 100644 async_image_generator.py
 create mode 100644 config.json
 create mode 100644 readme.md
 create mode 100644 remove_bg.py
 create mode 100644 requirements.txt

diff --git a/async_image_generator.py b/async_image_generator.py
new file mode 100644
index 0000000..93867e4
--- /dev/null
+++ b/async_image_generator.py
@@ -0,0 +1,190 @@
+"""
+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())
+
diff --git a/config.json b/config.json
new file mode 100644
index 0000000..594edb7
--- /dev/null
+++ b/config.json
@@ -0,0 +1,41 @@
+{"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"
+}
diff --git a/readme.md b/readme.md
new file mode 100644
index 0000000..e093166
--- /dev/null
+++ b/readme.md
@@ -0,0 +1,106 @@
+# Football Manager AI Face Generator
+
+## 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/)
+
+## Requirements
+
+- Python 3.7 or higher
+- Football Manager
+- NewGan or compatible face pack installer
+- DeepInfra API key
+
+## Installation
+
+1. Clone this repository or download the source files.
+2. Install required Python packages:
+
+`pip install -r requirements.txt`
+3. Generate and replace your DeepInfra API key in the `.env` file.
+
+
+## File Structure
+
+- `async_image_generator.py`: Main script for generating images
+- `config.json`: Configuration file for countries, facial characteristics, and prompts
+- `remove_bg.py`: Script to remove backgrounds from generated images
+- `.env`: File containing your DeepInfra API key
+
+## Usage
+
+### Generating Images
+
+Run the main script with optional arguments:
+
+`python async_image_generator.py [--n_per_country N] [--width W] [--height H] [--num_inference_steps S] [--model MODEL] [--resize]`
+
+Arguments:
+- `--n_per_country`: Number of images to generate per country group (default: 1)
+- `--width`: Width of generated images (default: 512)
+- `--height`: Height of generated images (default: 512)
+- `--num_inference_steps`: Number of inference steps (default: 1)
+- `--model`: Model to use, either "schnell" or "dev" (default: "schnell". "dev" is more EXPENSIVE)
+- `--resize`: If set, resizes images to 256x256
+
+Example:
+
+`python async_image_generator.py --n_per_country 2 --width 1024 --height 1024 --num_inference_steps 2 --model dev --resize`
+
+### Removing Backgrounds
+
+After generating images, you can remove backgrounds using:
+
+`python remove_bg.py [--directory DIR]`
+
+Arguments:
+- `--directory`: Path to the directory containing images (default: "generated_images")
+
+## Configuration
+
+Edit `config.json` to modify:
+- Country groups and countries within each group. The faces will be based on the countries.
+- Facial characteristics
+- Hair styles
+- The prompt template for image generation. I advise that you make minimal changes to the prompt.
+
+## Output
+
+Generated images are saved in the `generated_images` folder, organized by country group as it appears in any NewGan compatible facepack.
+
+```
+generated_images/
+├── African/
+├── Asian/
+├── Caucasian/
+├── Central European/
+├── EECA/
+├── Italmed/
+├── MENA/
+├── MESA/
+├── SAMed/
+├── Scandinavian/
+├── Seasian/
+├── South American/
+├── SpanMed/
+└── YugoGreek/
+```
+
+## Cost Calculation and User Confirmation
+
+Before beginning the image generation process, the script calculates the potential total cost based on the DeepInfra API pricing and the number of images to be generated (details in this link: https://deepinfra.com/black-forest-labs/FLUX-1-schnell/). It then displays this information to the user and asks for confirmation before proceeding. This allows users to make an informed decision about the cost before committing to the image generation process.
+
+
+## Troubleshooting
+
+- Ensure your `.env` file is correctly set up with a valid API key.
+- Check your internet connection if the script fails to connect to the API.
+- Make sure you have sufficient credits in your DeepInfra account.
+
+## Support
+
+For issues or questions, please open an issue on the GitHub repository.
+
+## License
+
+This project is open-source and available under the MIT License.
\ No newline at end of file
diff --git a/remove_bg.py b/remove_bg.py
new file mode 100644
index 0000000..6683b9b
--- /dev/null
+++ b/remove_bg.py
@@ -0,0 +1,75 @@
+"""
+Author: Emil Mirzayev
+
+This script removes the background from images in a specified directory and its subfolders.
+It processes JPG and JPEG files, converting them to PNG files with transparent backgrounds.
+"""
+
+import os
+import argparse
+from rembg import remove
+from PIL import Image
+import glob
+from tqdm import tqdm
+
+
+
+def remove_background(input_path, output_path):
+    """
+    Remove the background from a single image.
+
+    Args:
+    input_path (str): Path to the input image file.
+    output_path (str): Path where the processed image will be saved.
+    """
+    with Image.open(input_path) as img:
+        output = remove(img)
+        output.save(output_path)
+
+def process_directory(directory):
+    """
+    Process all JPG and JPEG images in the given directory and its subfolders.
+
+    Args:
+    directory (str): Path to the directory containing images.
+
+    Returns:
+    int: The number of images successfully processed.
+    """
+    processed_count = 0
+    
+    # Get the total number of files to process
+    total_files = sum(len(files) for _, _, files in os.walk(directory))
+    
+    # Create a progress bar
+    with tqdm(total=total_files, desc="Processing images", unit="image") as pbar:
+        for subdir, dirs, files in os.walk(directory):
+            for file in files:
+                if file.lower().endswith(('.jpg', '.jpeg', 'png')):
+                    input_path = os.path.join(subdir, file)
+                    output_filename = os.path.splitext(file)[0] + '.png'
+                    output_path = os.path.join(subdir, output_filename)
+                    
+                    try:
+                        remove_background(input_path, output_path)
+                        processed_count += 1
+                    except Exception as e:
+                        print(f"Error processing {input_path}: {str(e)}")
+                
+                # Update the progress bar
+                pbar.update(1)
+    
+    return processed_count
+def main():
+    """
+    Main function to parse command-line arguments and initiate the background removal process.
+    """
+    parser = argparse.ArgumentParser(description="Remove background from images in a directory and its subfolders")
+    parser.add_argument("--directory", type=str, default="generated_images", help="Path to the directory containing images. Defaults to `generated_images` folder in the same directory")
+    args = parser.parse_args()
+
+    total_processed = process_directory(args.directory)
+    print(f"\nTotal images processed: {total_processed}")
+
+if __name__ == "__main__":
+    main()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..327c4e1
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+aiohttp
+Pillow
+python-dotenv
+rembg
+tqdm
\ No newline at end of file