Spaces:
Sleeping
Sleeping
| 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}") |