Spaces:
Sleeping
Sleeping
File size: 4,166 Bytes
3b0ba50 |
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 |
import os
import zipfile
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from image_functions import generate_image, is_valid_image
from r2_uploader import upload_image_to_r2
from mongo_logger import create_log, update_log_status
def process_zip_and_generate_images(image_path, template_path, category, size, quality, user_prompt, sentiment):
try:
image_files = []
image_files.append((os.path.basename(template_path), template_path.name, os.path.basename(image_path.name), image_path.name))
results = process_image_files(image_files, category, size, quality, user_prompt, sentiment)
all_urls = [url for entry in results for url in entry["urls"]]
new_base = "https://adgen.statics.lookfinity.net/"
modified_urls = [new_base + url.rsplit('/', 1)[-1] for url in all_urls]
return modified_urls
except Exception as e:
log_error(f"Global error in zip processing: {e}")
return [], [f"Global error: {e}"]
def extract_zip_file(zip_path):
temp_dir = tempfile.TemporaryDirectory()
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(temp_dir.name)
return temp_dir
def get_valid_image_files(temp_dir):
valid_files = []
zip_contents = os.listdir(temp_dir.name)
for file in zip_contents:
file_path = os.path.join(temp_dir.name, file)
if is_valid_image(file) and "__MACOSX" not in file:
valid_files.append((file, file_path))
return valid_files
def process_image_files(image_files, category, size, quality, user_prompt, sentiment):
final_results = []
logs = []
with ThreadPoolExecutor(max_workers=5) as executor:
futures = []
for template_name, template_path, file_name, file_path in image_files:
log_id = create_log(category, file_name)
logs.append(f"Processing {file_name}")
futures.append(
executor.submit(
process_single_image,
template_path,
file_name,
file_path,
category,
size,
quality,
user_prompt,
sentiment,
log_id,
logs
)
)
for future in as_completed(futures):
result = future.result()
if result:
final_results.append(result)
return final_results
def process_single_image(template_path, file_name, file_path, category, size, quality, user_prompt, sentiment, log_id, logs):
try:
# base64_img = encode_image_to_base64(file_path)
# logs.append(f"Prompts generated for {file_name}")
#
# prompts = get_prompts(base64_img, category, user_prompt, sentiment, negative_prompt)
# logs.append(f"Generating images for {file_name}")
image_urls = generate_images_from_prompts(template_path, file_path, size, quality, category, sentiment, user_prompt)
status = "completed" if image_urls else "failed"
message = "Completed successfully" if image_urls else "No images generated"
update_log_status(log_id, status, urls=image_urls, message=message)
if image_urls:
return {
"file_name": file_name,
"urls": image_urls
}
return None
except Exception as e:
logs.append(f"Processing failed for {file_name}: {e}")
update_log_status(log_id, "failed", urls=[], message=str(e))
return None
def generate_images_from_prompts(template_path, file_path, size, quality, category, sentiment, user_prompt):
image_urls = []
for _ in range(10):
try:
img_url = generate_image(template_path, file_path, size, quality, category, sentiment, user_prompt)
s3_url = upload_image_to_r2(img_url)
image_urls.append(s3_url)
except Exception as e:
log_error(f"Image generation failed: {e}")
return image_urls
def log_error(message):
print(f"ERROR: {message}") |