Update generator_function/multimodel_image_processor.py

generator_function/multimodel_image_processor.py

@@ -36,6 +36,7 @@ def process_zip_and_generate_images_multimodel(
     uid: str,
     selected_model: str = "gpt_default",
     model_params: Optional[dict] = None,
+    private_mode: bool = False,
 ) -> List[str]:
     """Enhanced image processor that supports both GPT and multimodel approaches"""
     num_images = 1 if demo_mode else num_images
@@ -48,7 +49,7 @@ def process_zip_and_generate_images_multimodel(
 
         results = process_image_files_multimodel(
             image_files, category, size, quality, user_prompt, sentiment, platform, 
-            num_images, blur, uid, selected_model, model_params
+            num_images, blur, uid, selected_model, model_params, private_mode
         )
         all_urls = [url for entry in results for url in entry["urls"]]
         seen, deduped = set(), []
@@ -82,14 +83,14 @@ def get_valid_image_files(temp_dir: tempfile.TemporaryDirectory) -> List[Tuple[s
 
 def process_image_files_multimodel(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, selected_model: str, model_params: Optional[dict]) -> List[dict]:
+    uid: str, selected_model: str, model_params: Optional[dict], private_mode: bool) -> List[dict]:
     """Process image files with multimodel support"""
     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, 
@@ -112,7 +113,7 @@ def process_image_files_multimodel(image_files: List[Tuple[str, str]], category:
                 executor.submit(
                     process_single_image_multimodel,
                     file_name, file_path, category, size, quality, user_prompt, sentiment, 
-                    platform, num_images, blur, job_id, selected_model, model_params,
+                    platform, num_images, blur, job_id, selected_model, model_params, private_mode,
                 )
             )
         for future in as_completed(futures):
@@ -125,14 +126,14 @@ def process_image_files_multimodel(image_files: List[Tuple[str, str]], category:
 
 def process_single_image_multimodel(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], selected_model: str, model_params: Optional[dict]) -> Optional[dict]:
+    job_id: Optional[str], selected_model: str, model_params: Optional[dict], private_mode: bool) -> Optional[dict]:
     """Process single image with multimodel support"""
     try:
         image_urls = generate_images_from_prompts_multimodel(
             file_path, size, quality, category, sentiment, user_prompt, platform, 
-            num_images, blur, selected_model, model_params
+            num_images, blur, selected_model, model_params, 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:
@@ -142,7 +143,7 @@ def process_single_image_multimodel(file_name: str, file_path: str, category: st
         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:
@@ -152,6 +153,7 @@ def process_single_image_multimodel(file_name: str, file_path: str, category: st
 def generate_images_from_prompts_multimodel(
     file_path: str, size: str, quality: str, category: str, sentiment: str, user_prompt: str,
     platform: str, num_images: int, blur: bool, selected_model: str, model_params: Optional[dict],
+    private_mode: bool,
 ) -> List[str]:
     """Generate images using either GPT or multimodel approach"""
     image_urls: List[str] = []
@@ -170,6 +172,8 @@ def generate_images_from_prompts_multimodel(
             
             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: