gpu support

comfy_fm_newgen.py

@@ -20,7 +20,7 @@ from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
 
 logging.config.dictConfig(LOGGING_CONFIG)
 
-cut = 500
+cut = 200
 
 # Load user configurations
 user_config = configparser.ConfigParser()

config.json

@@ -88,5 +88,5 @@
     "3": "Long",
     "4": "Bald"
   },
-  "prompt": "Ultra-realistic close-up headshot of a {skin_tone} skinned male soccer player with a white background looking at the camera. The player is {age} years old, from {country}, with {facial_characteristics} and {hair} hair. He is facing the camera with a confident expression, wearing a soccer jersey. The lighting is natural and soft, emphasizing facial features and skin texture"
+  "prompt": "Ultra-realistic close-up headshot of a {skin_tone} skinned male soccer player with a plain background looking at the camera. The player is {age} years old, from {country}, with {facial_characteristics} and {hair} hair. He is facing the camera with a confident expression, wearing a soccer jersey. The lighting is natural and soft, emphasizing facial features and skin texture"
 }

lib/remove_bg.py

@@ -1,61 +1,70 @@
-import os
 from rembg import remove
 from PIL import Image
 from tqdm import tqdm
+import os
 from concurrent.futures import ThreadPoolExecutor
 
-def process_image(input_path, output_path):
+import onnxruntime as ort
-    """
-    Process a single image: remove its background and save the result.
-    
-    Args:
-    input_path (str): Path to the input image.
-    output_path (str): Path to save the processed image.
-    
-    Returns:
-    bool: True if the image was successfully processed, False otherwise.
-    """
-    try:
-        with Image.open(input_path) as img:
-            output = remove(img)
-            output.save(output_path)
-        return True
-    except Exception as e:
-        print(f"Error processing {input_path}: {str(e)}")
-        return False
 
-def remove_bg_from_files_in_dir(directory, max_workers=2):
+# Suppress ONNX Runtime logging to show only critical errors
+ort.set_default_logger_severity(3)  # 0 = verbose, 1 = info, 2 = warning, 3 = error, 4 = fatal
+
+
+def process_images_in_batch(batch):
     """
-    Process all JPG, JPEG, and PNG images in the given directory and its subfolders using parallel processing.
+    Process a batch of images: remove their backgrounds and save the results.
+
+    Args:
+    batch (list): List of tuples (input_path, output_path).
+
+    Returns:
+    int: Number of images successfully processed in this batch.
+    """
+    success_count = 0
+    for input_path, output_path in batch:
+        try:
+            with Image.open(input_path) as img:
+                output = remove(img)  # This will use GPU if ONNX Runtime is GPU-enabled
+                output.save(output_path)
+            success_count += 1
+        except Exception as e:
+            print(f"Error processing {input_path}: {str(e)}")
+    return success_count
+
+def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=5):
+    """
+    Process all JPG, JPEG, and PNG images in the given directory and its subfolders using parallel processing and GPU.
 
     Args:
     directory (str): Path to the directory containing images.
     max_workers (int): Maximum number of threads to use for parallel processing.
-    
+    batch_size (int): Number of images to process per batch.
+
     Returns:
     int: The number of images successfully processed.
     """
-    processed_count = 0
     files_to_process = []
 
     # Gather all the image files to process
     for subdir, dirs, files in os.walk(directory):
         for file in files:
-            if file.lower().endswith(('.jpg', '.jpeg', 'png')):
+            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)
                 files_to_process.append((input_path, output_path))
 
-    # Use ThreadPoolExecutor to process images in parallel
+    processed_count = 0
+
+    # Divide files into batches
+    batches = [files_to_process[i:i + batch_size] for i in range(0, len(files_to_process), batch_size)]
+
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        # Use tqdm to show progress
         with tqdm(total=len(files_to_process), desc="Processing images", unit="image") as pbar:
-            futures = {executor.submit(process_image, input_path, output_path): (input_path, output_path) for input_path, output_path in files_to_process}
+            futures = {executor.submit(process_images_in_batch, batch): batch for batch in batches}
 
             for future in futures:
-                if future.result():
+                processed_count += future.result()
-                    processed_count += 1
+                pbar.update(len(futures[future]))
-                pbar.update(1)
 
     return processed_count