Spaces:
Sleeping
Sleeping
File size: 10,735 Bytes
b2cbfed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 |
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...===") |