latest working commit

.vscode/launch.json

@@ -15,7 +15,7 @@
                 "--num_inference_steps",
                 "6",
                 "--rtf_file",
-                "./2030New.rtf"
+                "./NewGen.rtf"
             ]
         }
     ]

comfy_fm_newgen.py

@@ -5,14 +5,18 @@ import json
 import configparser
 import pycountry
 import inflect
-from lib.rtf_parser import RTF_Parser
 import logging
 import sys
+
+from tqdm import tqdm
+from lib.rtf_parser import RTF_Parser
+from lib.remove_bg import remove_bg_from_files_in_dir
+from lib.generate_xml import create_config_xml
+from lib.resize_images import resize_images
 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")
@@ -20,8 +24,10 @@ user_config.read("config.cfg")
 rtf = RTF_Parser()
 p = inflect.engine()
 
+output_folder = user_config["general"]["output_dir"]
 
-def generate_image(uid,comfy_prompt,app_config,model):
+
+def generate_image(uid, comfy_prompt, model):
     """Generate an image using the Comfy API."""
     url = user_config["general"]["url"]
 
@@ -42,6 +48,7 @@ def generate_image(uid,comfy_prompt,app_config,model):
         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:
@@ -54,10 +61,9 @@ def get_country_name_from_code(code):
 def generate_prompts_for_players(players, app_config):
     """Generate images for a specific player and configuration."""
     prompts = []
-    for player in players[:100]:
+    for player in players:
         print(f"\nGenerating image for {player[0]} - {player[8]}")
-        folder_name = f"generated_images/"
+        os.makedirs(output_folder, exist_ok=True)
-        os.makedirs(folder_name, exist_ok=True)
 
         country = get_country_name_from_code(player[1])
         facial_characteristics = random.choice(app_config["facial_characteristics"])
@@ -80,6 +86,7 @@ def generate_prompts_for_players(players, app_config):
         prompts.append(prompt)
     return prompts
 
+
 def main():
     """Main function for generating images."""
     parser = argparse.ArgumentParser(description="Generate images for country groups")
@@ -101,7 +108,7 @@ def main():
         raise Exception("Please pass in a RTF file as --rtf_file")
 
     # Parse the RTF file
-    rtf_file = rtf.parse_rtf(args.rtf_file)
+    rtf_file = rtf.parse_rtf(args.rtf_file)[:5]
 
     # Extract unique values
     hair_length = list(set(item[5] for item in rtf_file))
@@ -111,19 +118,21 @@ def main():
     # Load configurations
     with open("config.json", "r") as f:
         app_config = json.load(f)
-
+    print(f"{len(rtf_file)} images will be generated")
     prompts = generate_prompts_for_players(rtf_file, app_config)
-    for prompt in prompts:
+    for prompt in tqdm(prompts, desc="Generating Images"):
         uid = prompt.split(":")[0]
         comfy_prompt = prompt.split(":")[1]
         generate_image(
             uid,
             comfy_prompt,
-            app_config,
             user_config["general"]["model"],
         )
+    remove_bg_from_files_in_dir(output_folder)
+    resize_images(output_folder)
+    create_config_xml(output_folder)
+    print("\nImage generation complete for players in RTF file.")
 
-    print("\nImage generation complete for all country groups.")
 
 if __name__ == "__main__":
     main()

config.cfg.sample

@@ -1,3 +1,4 @@
 [general]
 url = 
 model = realisticVisionV60B1_v51HyperVAE.safetensors
+output_dir = ./generated_images/

lib/generate_xml.py

@@ -0,0 +1,30 @@
+import os
+import xml.etree.ElementTree as ET
+
+def create_config_xml(folder_path):
+    # Define the root element
+    root = ET.Element("record")
+
+    # Add resource manager options
+    ET.SubElement(root, "boolean", id="preload", value="false")
+    ET.SubElement(root, "boolean", id="amap", value="false")
+
+    # Add the maps section
+    maps = ET.SubElement(root, "list", id="maps")
+
+    # Iterate through files in the folder
+    for filename in os.listdir(folder_path):
+        if filename.endswith(".png"):
+            # Extract the UID from the filename (remove the extension)
+            uid = os.path.splitext(filename)[0]
+
+            # Create a record for each file
+            ET.SubElement(maps, "record", from_=uid, to=f"graphics/pictures/person/r-{uid}/portrait")
+
+    # Create the XML tree
+    tree = ET.ElementTree(root)
+
+    # Save the XML file
+    output_path = os.path.join(folder_path, "config.xml")
+    tree.write(output_path, encoding="utf-8", xml_declaration=True)
+    print(f"Config XML created at: {output_path}")

lib/remove_bg.py

@@ -0,0 +1,41 @@
+import os
+from rembg import remove
+from PIL import Image
+from tqdm import tqdm
+
+def remove_bg_from_files_in_dir(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:
+                        with Image.open(input_path) as img:
+                            output = remove(img)
+                            output.save(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

lib/resize_images.py

@@ -0,0 +1,17 @@
+import os
+from PIL import Image
+from tqdm import tqdm
+
+
+def resize_images(folder_path):
+    # Loop through all files in the folder
+    for filename in tqdm(os.listdir(folder_path), desc="Resizing Images"):
+        # Check if the file is a PNG image
+        if filename.endswith('.png'):
+            # Full path to the image file
+            image_path = os.path.join(folder_path, filename)
+            # Open the image file
+            with Image.open(image_path) as img:
+                # Resize the image to 256x256
+                img_resized = img.resize((256, 256))
+                img_resized.save(os.path.join(folder_path, filename))