Update generator_function/image_processor.py

generator_function/image_processor.py

@@ -1,4 +1,4 @@
-import os, zipfile, tempfile, logging
+import os, zipfile, tempfile, logging, base64
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Tuple, Optional
 
@@ -29,6 +29,7 @@ def process_zip_and_generate_images(
     existing_images: Optional[List[str]],
     blur: bool,
     uid: str,
+    private_mode: bool = False,
 ) -> List[str]:
     num_images = 1 if demo_mode else num_images
     try:
@@ -39,7 +40,7 @@ def process_zip_and_generate_images(
             image_files = [(os.path.basename(zip_path), zip_path)]
 
         results = process_image_files(
-            image_files, category, size, quality, user_prompt, sentiment, platform, num_images, blur, uid
+            image_files, category, size, quality, user_prompt, sentiment, platform, num_images, blur, uid, private_mode
         )
         all_urls = [url for entry in results for url in entry["urls"]]
         seen, deduped = set(), []
@@ -73,13 +74,14 @@ def get_valid_image_files(temp_dir: tempfile.TemporaryDirectory) -> List[Tuple[s
 
 def process_image_files(image_files: List[Tuple[str, str]],category: str,size: str,
     quality: str,user_prompt: str,sentiment: str,platform: str,num_images: int,blur: bool,uid: str,
+    private_mode: bool,
 ) -> List[dict]:
     final_results: List[dict] = []
     with ThreadPoolExecutor(max_workers=5) as executor:
         futures = []
         for file_name, file_path in image_files:
             job_id: Optional[str] = None
-            if COL is not None:
+            if COL is not None and not private_mode:
                 try:
                     settings = {"size": size,"quality": quality,"sentiment": sentiment,"platform": platform,"num_images": num_images,"blur": bool(blur)}
                     inputs = {"file_name": file_name, "mode": "img_or_zip"}
@@ -98,7 +100,7 @@ def process_image_files(image_files: List[Tuple[str, str]],category: str,size: s
             futures.append(
                 executor.submit(
                     process_single_image,
-                    file_name, file_path, category, size, quality, user_prompt, sentiment, platform, num_images, blur, job_id,
+                    file_name, file_path, category, size, quality, user_prompt, sentiment, platform, num_images, blur, job_id, private_mode,
                 )
             )
         for future in as_completed(futures):
@@ -110,12 +112,12 @@ def process_image_files(image_files: List[Tuple[str, str]],category: str,size: s
     return final_results
 
 def process_single_image(file_name: str,file_path: str,category: str,size: str,quality: str,user_prompt: str,sentiment: str,
-    platform: str,num_images: int,blur: bool,job_id: Optional[str],) -> Optional[dict]:
+    platform: str,num_images: int,blur: bool,job_id: Optional[str], private_mode: bool) -> Optional[dict]:
     try:
         image_urls = generate_images_from_prompts(
-            file_path, size, quality, category, sentiment, user_prompt, platform, num_images, blur
+            file_path, size, quality, category, sentiment, user_prompt, platform, num_images, blur, private_mode
         )
-        if COL is not None and job_id:
+        if COL is not None and job_id and not private_mode:
             try:
                 finish_job(COL, job_id, status=("completed" if image_urls else "failed"), outputs_urls=image_urls)
             except Exception:
@@ -125,7 +127,7 @@ def process_single_image(file_name: str,file_path: str,category: str,size: str,q
         return None
     except Exception as e:
         logger.error(f"Processing failed for {file_name}: {e}")
-        if COL is not None and job_id:
+        if COL is not None and job_id and not private_mode:
             try:
                 finish_job(COL, job_id, status="failed", outputs_urls=[])
             except Exception:
@@ -134,7 +136,7 @@ def process_single_image(file_name: str,file_path: str,category: str,size: str,q
 
 def generate_images_from_prompts(
     file_path: str, size: str, quality: str, category: str, sentiment: str, user_prompt: str,
-    platform: str, num_images: int, blur: bool,
+    platform: str, num_images: int, blur: bool, private_mode: bool,
 ) -> List[str]:
     image_urls: List[str] = []
 
@@ -143,6 +145,8 @@ def generate_images_from_prompts(
             image_bytes = generate_image(file_path, size, quality, category, sentiment, user_prompt, platform, blur, i)
             if not image_bytes: return None
             image_with_metadata = meta_data_helper_function(image_bytes)
+            if private_mode:
+                return "data:image/png;base64," + base64.b64encode(image_with_metadata).decode("utf-8")
             s3_url = upload_image_to_r2(image_with_metadata)
             return s3_url
         except Exception as e: