Update generator_function/image_function.py

generator_function/image_function.py

@@ -147,8 +147,8 @@ def _fetch(url: Union[str, Any]) -> Optional[bytes]:
             time.sleep(1)
     return None
 
-def _process_one(args: Tuple[str, str, str, int]) -> Dict[str, Any]:
-    model_key, prompt, aspect_ratio, idx = args
+def _process_one(args: Tuple[str, str, str, int, bool]) -> Dict[str, Any]:
+    model_key, prompt, aspect_ratio, idx, private_mode = args
     out = {"index": idx, "success": False, "source_url": None, "r2_url": None, "error": None}
     try:
         urls = _generate_one(model_key, prompt, aspect_ratio)
@@ -160,22 +160,27 @@ def _process_one(args: Tuple[str, str, str, int]) -> Dict[str, Any]:
         if not b:
             out["error"] = "Fetch failed"; return out
         image_with_metadata = meta_data_helper_function(b)
-        r2 = _upload_to_r2(image_with_metadata)
-        if r2:
-            out["r2_url"] = r2; out["success"] = True
+        if private_mode:
+            data_uri = "data:image/png;base64," + base64.b64encode(image_with_metadata).decode("utf-8")
+            out["r2_url"] = data_uri
+            out["success"] = True
         else:
-            out["error"] = "Upload to R2 failed"
+            r2 = _upload_to_r2(image_with_metadata)
+            if r2:
+                out["r2_url"] = r2; out["success"] = True
+            else:
+                out["error"] = "Upload to R2 failed"
     except Exception as e:
         out["error"] = str(e)
     return out
 
-def _generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str, num_images: int) -> Tuple[List[str], List[str], List[str]]:
+def _generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str, num_images: int, *, private_mode: bool = False) -> Tuple[List[str], List[str], List[str]]:
     if num_images == 1:
-        res = _process_one((model_key, prompt, aspect_ratio, 0))
+        res = _process_one((model_key, prompt, aspect_ratio, 0, private_mode))
         if res["success"]:
             return [res["r2_url"]], [res["source_url"]], []
         return [], [], [res["error"] or "Generation failed"]
-    args = [(model_key, prompt, aspect_ratio, i) for i in range(num_images)]
+    args = [(model_key, prompt, aspect_ratio, i, private_mode) for i in range(num_images)]
     r2, src, errs = [], [], []
     with ThreadPoolExecutor(max_workers=min(MAX_WORKERS, num_images)) as ex:
         for fut in as_completed({ex.submit(_process_one, a): a[3] for a in args}):
@@ -194,9 +199,9 @@ def _generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str, nu
 
 
 
-def generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str, num_images: int) -> Tuple[List[str], List[str], List[str]]:
+def generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str, num_images: int, *, private_mode: bool = False) -> Tuple[List[str], List[str], List[str]]:
     """Back-compat public export used by background tasks."""
-    return _generate_images_parallel(model_key, aspect_ratio, prompt, num_images)
+    return _generate_images_parallel(model_key, aspect_ratio, prompt, num_images, private_mode=private_mode)
 
 def handle_image_generation_optimized(
     *,
@@ -208,6 +213,7 @@ def handle_image_generation_optimized(
     category: Optional[str] = None,
     platform: Optional[str] = None,
     uid:str,
+    private_mode: bool = False,
 
 ):
     """
@@ -223,7 +229,7 @@ def handle_image_generation_optimized(
 
     created_by = uid
 
-    results_col = get_results_collection()
+    results_col = None if private_mode else get_results_collection()
     db_job_id = None
     if results_col is not None:
         try:
@@ -248,8 +254,14 @@ def handle_image_generation_optimized(
             st.info(f"Generating {num_images} image(s)")
         progress.progress(10, text="Running...")
 
-        r2_urls, source_urls, errors = _generate_images_parallel(model_key, aspect_ratio, prompt.strip(), num_images)
-        urls = r2_urls or source_urls
+        r2_urls, source_urls, errors = _generate_images_parallel(
+            model_key,
+            aspect_ratio,
+            prompt.strip(),
+            num_images,
+            private_mode=private_mode,
+        )
+        urls = r2_urls if private_mode else (r2_urls or source_urls)
 
         if results_col is not None and db_job_id:
             try:
@@ -268,7 +280,12 @@ def handle_image_generation_optimized(
 
         if urls:
             with status.container():
-                st.success(f"Generated {len(urls)} image(s) in {took:.1f}s. Job ID: {db_job_id or 'N/A'}")
+                message = f"Generated {len(urls)} image(s) in {took:.1f}s."
+                if not private_mode:
+                    message += f" Job ID: {db_job_id or 'N/A'}"
+                else:
+                    message += " Private mode: results stay local to this session."
+                st.success(message)
 
             cols = st.columns(min(4, len(urls)) or 1)
             image_bytes_list = []
@@ -276,7 +293,14 @@ def handle_image_generation_optimized(
             for i, u in enumerate(urls):
                 with cols[i % len(cols)]:
                     try:
-                        b = _fetch(u)
+                        if isinstance(u, str) and u.startswith("data:image"):
+                            try:
+                                _, encoded = u.split(",", 1)
+                                b = base64.b64decode(encoded)
+                            except Exception:
+                                b = None
+                        else:
+                            b = _fetch(u)
                         if b is None:
                             st.error("Failed to load image")
                             continue