diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index 96112ff..abf5ef2 100644 --- a/comfy_fm_newgen.py +++ b/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() diff --git a/config.json b/config.json index 335db91..25f4d8a 100644 --- a/config.json +++ b/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" } \ No newline at end of file diff --git a/lib/remove_bg.py b/lib/remove_bg.py index d8bd4f0..4199bf8 100644 --- a/lib/remove_bg.py +++ b/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): - """ - 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 +import onnxruntime as ort -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 += 1 - pbar.update(1) + processed_count += future.result() + pbar.update(len(futures[future])) return processed_count