import os import random import csv import gc import glob from datetime import datetime import time from pathlib import Path from style_template import style_list from PIL import Image, ImageOps # Default Configuration variables INPUT_FOLDER_NAME = 'examples' OUTPUT_FOLDER_NAME = 'generated_images' LOG_FILENAME = 'generation_log.csv' logfile_path = os.path.join(os.getcwd(), LOG_FILENAME) PROMPT = "human, sharp focus" NEGATIVE_PROMPT = "(blurry, blur, text, abstract, glitch, lowres, low quality, worst quality:1.2), (text:1.2), watermark, painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured" IDENTITYNET_STRENGTH_RATIO_RANGE = (1.0, 1.5) ADAPTER_STRENGTH_RATIO_RANGE = (0.7, 1.0) NUM_INFERENCE_STEPS_RANGE = (40, 60) GUIDANCE_SCALE_RANGE = (7.0, 12.0) MAX_SIDE = 1280 MIN_SIDE = 1024 NUMBER_OF_LOOPS = 1 # Dynamically create the STYLES list from imported style_list STYLES = [style["name"] for style in style_list] USE_RANDOM_STYLE = False def choose_random_style(): return random.choice(STYLES) def get_random_image_file(input_folder): valid_extensions = [".jpg", ".jpeg", ".png"] files = [file for file in Path(input_folder).glob("*") if file.suffix.lower() in valid_extensions] if not files: raise FileNotFoundError(f"No images found in directory {input_folder}") return str(random.choice(files)) def resize_and_pad_image(image_path, max_side, min_side, pad_color=(255, 255, 255)): # Open an image using PIL image = Image.open(image_path) # Calculate the scale and new size ratio = min(min_side / min(image.size), max_side / max(image.size)) new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio)) # Resize the image image = image.resize(new_size, Image.BILINEAR) # Calculate padding delta_w = max_side - new_size[0] delta_h = max_side - new_size[1] # Pad the resized image to make it square padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2)) image = ImageOps.expand(image, padding, pad_color) return image def log_to_csv(logfile_path, image_name, new_file_name='Unknown', identitynet_strength_ratio=0.0, adapter_strength_ratio=0.0, num_inference_steps=0, guidance_scale=0.0, seed=0, success=True, error_message='', style_name="", prompt="", negative_prompt="", time_taken=0.0, current_timestamp=""): os.makedirs(os.path.dirname(logfile_path), exist_ok=True) file_exists = os.path.isfile(logfile_path) with open(logfile_path, 'a', newline='', encoding='utf-8') as csvfile: fieldnames = ['image_name', 'new_file_name', 'identitynet_strength_ratio', 'adapter_strength_ratio', 'num_inference_steps', 'guidance_scale', 'seed', 'success', 'error_message', 'style_name', 'prompt', 'negative_prompt', 'time_taken', 'current_timestamp'] writer = csv.DictWriter(csvfile, fieldnames=fieldnames) if not file_exists: writer.writeheader() writer.writerow({ 'image_name': image_name, 'new_file_name': new_file_name, 'identitynet_strength_ratio': identitynet_strength_ratio, 'adapter_strength_ratio': adapter_strength_ratio, 'num_inference_steps': num_inference_steps, 'guidance_scale': guidance_scale, 'seed': seed, 'success': success, 'error_message': error_message, 'style_name': style_name, 'prompt': prompt, 'negative_prompt': negative_prompt, 'time_taken': time_taken, 'current_timestamp': current_timestamp }) def initial_image(generate_image_func): overall_start_time = time.time() total_time_taken = 0.0 # Initialize a counter for processed images at the beginning of the function processed_images_count = 0 # List all image files in the `INPUT_FOLDER_NAME` image_files = glob.glob(f'{INPUT_FOLDER_NAME}/*.png') + \ glob.glob(f'{INPUT_FOLDER_NAME}/*.jpg') + \ glob.glob(f'{INPUT_FOLDER_NAME}/*.jpeg') # Check if we found any images if not image_files: raise FileNotFoundError(f"No images found in directory {INPUT_FOLDER_NAME}") # Print the count of detected image files print(f"Processing a total of {len(image_files)} image(s) in '{INPUT_FOLDER_NAME}'") # Shuffle the image files randomly random.shuffle(image_files) total_images = len(image_files) # Get the total number of images to process for loop in range(NUMBER_OF_LOOPS): print(f"Starting loop {loop+1} of {NUMBER_OF_LOOPS}") for image_number, face_image_path in enumerate(image_files, start=1): loop_start_time = datetime.now() face_image = [face_image_path] basename = os.path.basename(face_image_path) processed_images_count += 1 # Resize and pad the image before processing processed_image = resize_and_pad_image( image_path=face_image_path, max_side=MAX_SIDE, min_side=MIN_SIDE ) if USE_RANDOM_STYLE: style_name = choose_random_style() else: style_name = "(No style)" identitynet_strength_ratio = random.uniform(*IDENTITYNET_STRENGTH_RATIO_RANGE) adapter_strength_ratio = random.uniform(*ADAPTER_STRENGTH_RATIO_RANGE) num_inference_steps = random.randint(*NUM_INFERENCE_STEPS_RANGE) guidance_scale = random.uniform(*GUIDANCE_SCALE_RANGE) seed = random.randint(0, 2**32 - 1) # Print settings for the current image BEFORE processing it print_generation_settings(basename, style_name, identitynet_strength_ratio, adapter_strength_ratio, num_inference_steps, guidance_scale, seed, image_number, total_images) # Here, the generate_image_func is supposedly called and image processing happens _, _, generated_file_paths = generate_image_func( face_image=face_image, pose_image=None, prompt=PROMPT, negative_prompt=NEGATIVE_PROMPT, style_name=style_name, enhance_face_region=True, num_steps=num_inference_steps, identitynet_strength_ratio=identitynet_strength_ratio, adapter_strength_ratio=adapter_strength_ratio, guidance_scale=guidance_scale, seed=seed ) loop_end_time = datetime.now() loop_time_taken = (loop_end_time - loop_start_time).total_seconds() # Immediately print the time taken and current time. print(f"Time taken to process image: {loop_time_taken:.2f} seconds") # Update the total time taken with this image's processing time total_time_taken += loop_time_taken # Calculate the average time taken per image average_time_per_image = total_time_taken / image_number current_timestamp = loop_end_time.strftime("%Y-%m-%d %H:%M:%S") # Current time after processing print(f"Current timestamp: {current_timestamp}") # Calculate estimated remaining time considering the images left in this loop and the additional loops remaining_images_this_loop = total_images - image_number remaining_images_in_additional_loops = (NUMBER_OF_LOOPS - (loop + 1)) * total_images total_remaining_images = remaining_images_this_loop + remaining_images_in_additional_loops estimated_time_remaining = average_time_per_image * total_remaining_images # Display the estimated time remaining including remaining loops print(f"Estimated time remaining (including loops): {estimated_time_remaining // 60:.0f} minutes, {estimated_time_remaining % 60:.0f} seconds") # Display the overall average time per image in seconds print(f"Overall average time per image: {average_time_per_image:.2f} seconds") # Display the total number of remaining images to process including looping print(f"Total remaining images to process (including loops): {total_remaining_images}") success = True # Assuming generation was successful. error_message = "" # Assuming no error. # Log to CSV after the image generation. for generated_file_path in generated_file_paths: new_file_name = os.path.basename(generated_file_path) log_to_csv(logfile_path, basename, new_file_name, identitynet_strength_ratio, adapter_strength_ratio, num_inference_steps, guidance_scale, seed, success, error_message, style_name, PROMPT, NEGATIVE_PROMPT, loop_time_taken, current_timestamp) del generated_file_paths # Explicitly delete large variables gc.collect() # Call garbage collection # At the end of the initial_image() function, add: total_elapsed_time = time.time() - overall_start_time print("\n===FINAL SUMMARY===") print(f"Total loops completed: {NUMBER_OF_LOOPS}") print(f"Total images processed per loop: {len(image_files)}") print(f"Overall total images processed: {NUMBER_OF_LOOPS * len(image_files)}") # Multiplied by the number of loops print(f"Overall total time: {total_elapsed_time / 60:.2f} minutes") def print_generation_settings(basename, style_name, identitynet_strength_ratio, adapter_strength_ratio, num_inference_steps, guidance_scale, seed, image_number, total_images): print("===IMAGE GENERATION DATA SUMMARY===") # Print settings for the current image print(f"- Image {image_number} of {total_images}\n" f"- Filename: {basename}\n" f"- Style: {style_name}\n" f"- IdentityNet strength ratio: {identitynet_strength_ratio:0.2f}\n" f"- Adapter strength ratio: {adapter_strength_ratio:0.2f}\n" f"- Number of inference steps: {num_inference_steps}\n" f"- Guidance scale: {guidance_scale:0.2f}\n" f"- Seed: {seed}\n" f"- Input folder name: {INPUT_FOLDER_NAME}\n" f"- Output folder name: {OUTPUT_FOLDER_NAME}\n" f"- Prompt: {PROMPT}\n" f"- Negative prompt: {NEGATIVE_PROMPT}\n" f"- Number of loops: {NUMBER_OF_LOOPS}\n" f"- Use random style: {USE_RANDOM_STYLE}\n") print("===DEFINING COMPLETE, GENERATING IMAGE...===")