rework with module

.vscode/launch.json

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

comfy_fm_newgen.py

@@ -1,114 +1,49 @@
-"""
+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 random
 import json
+import configparser
 import pycountry
 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.read('config.cfg')
+user_config.read("config.cfg")
 
 rtf = RTF_Parser()
 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 = 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))
-
-#         # 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.")
     
+    # 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(
+def generate_prompts_for_players(players, app_config):
-    session,
+    """Generate images for a specific player and configuration."""
-    player,
+    prompts = []
-    config,
+    for player in players[:100]:
-    model,
+        print(f"\nGenerating image for {player[0]} - {player[8]}")
-):
+        folder_name = f"generated_images/"
-    print(f"\nGenerating image for {player[0]} - {player[8]}")
+        os.makedirs(folder_name, exist_ok=True)
-    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"])
+        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])
-        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=facial_characteristics or "no facial hair",
-                facial_characteristics if facial_characteristics else "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(
+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.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()
+
     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
+    # Extract unique values
-    hair_length = list(
+    hair_length = list(set(item[5] for item in rtf_file))
-        set(item[pos] for item in rtf_file for pos in [5])
+    hair_colour = list(set(item[6] for item in rtf_file))
-    )
+    skin_tone = list(set(item[7] for item in rtf_file))
-    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])
-        )
 
+    # Load configurations
     with open("config.json", "r") as 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(f"Total images generated: {generated_images}")
-
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()

readme.md

@@ -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/).
 
 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. -->

workflow.json

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

workflow_api.json

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