Update app_pages/image_generation.py

app_pages/image_generation.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import zipfile, requests, logging, tempfile, shutil, os
+import zipfile, requests, logging, tempfile, shutil, os, base64
 import streamlit as st
 from generator_function.image_processor import process_zip_and_generate_images
 from generator_function.multimodel_image_processor import process_zip_and_generate_images_multimodel
@@ -12,24 +12,53 @@ logger = logging.getLogger(__name__)
 def _zip_gallery_images(gallery_items):
     if not gallery_items:
         return None
-    image_urls = [item[0] if isinstance(item, tuple) else item for item in gallery_items]
     temp_dir = tempfile.mkdtemp()
     try:
         zip_path = tempfile.NamedTemporaryFile(delete=False, suffix=".zip").name
-        for i, url in enumerate(image_urls):
+        for i, item in enumerate(gallery_items):
+            filename = None
+            url = None
+            data_bytes = None
+
+            if isinstance(item, tuple):
+                if len(item) == 2 and isinstance(item[1], (bytes, bytearray)):
+                    filename = item[0]
+                    data_bytes = item[1]
+                else:
+                    url = item[0]
+            elif isinstance(item, str):
+                if item.startswith("data:image"):
+                    try:
+                        header, encoded = item.split(",", 1)
+                        mime = header.split(";")[0].split("/")[-1] or "png"
+                        data_bytes = base64.b64decode(encoded)
+                        filename = f"image_{i}.{mime}"
+                    except Exception:
+                        logger.error("Failed to decode private image data for zip bundle.")
+                        continue
+                else:
+                    url = item
+            else:
+                continue
+
             try:
-                ext = url.split('?')[0].split('.')[-1]
-                ext = ext if ext and len(ext) <= 5 else "png"
-                file_path = os.path.join(temp_dir, f"image_{i}.{ext}")
-                if url.startswith(("http://", "https://")):
-                    resp = requests.get(url, timeout=10)
-                    resp.raise_for_status()
+                if data_bytes is not None:
+                    file_path = os.path.join(temp_dir, filename or f"image_{i}.png")
                     with open(file_path, "wb") as f:
-                        f.write(resp.content)
-                elif os.path.exists(url):
-                    shutil.copy(url, file_path)
+                        f.write(data_bytes)
+                elif url:
+                    ext = url.split('?')[0].split('.')[-1]
+                    ext = ext if ext and len(ext) <= 5 else "png"
+                    file_path = os.path.join(temp_dir, f"image_{i}.{ext}")
+                    if url.startswith(("http://", "https://")):
+                        resp = requests.get(url, timeout=10)
+                        resp.raise_for_status()
+                        with open(file_path, "wb") as f:
+                            f.write(resp.content)
+                    elif os.path.exists(url):
+                        shutil.copy(url, file_path)
             except Exception as e:
-                logger.error(f"Error processing image {url}: {e}")
+                logger.error(f"Error processing image {url or filename}: {e}")
         with zipfile.ZipFile(zip_path, "w") as zipf:
             for file_name in os.listdir(temp_dir):
                 zipf.write(os.path.join(temp_dir, file_name), arcname=file_name)
@@ -42,6 +71,7 @@ def _zip_gallery_images(gallery_items):
 
 
 def render_bulk_image_generator(uid: str, prefix: str = "ig_img"):
+    private_mode = st.session_state.get("private_mode_enabled", False)
     zip_file = st.file_uploader(
         "Upload Zip or Single File",
         type=["zip", "png", "jpg", "jpeg"],
@@ -136,6 +166,9 @@ def render_bulk_image_generator(uid: str, prefix: str = "ig_img"):
     gallery = st.session_state.setdefault(gallery_key, [])
 
 
+    if private_mode:
+        st.info("Private mode is active. Generated variations stay only in this session and are not uploaded to storage or the AI Library.")
+
     if demo_btn or gen_all_btn:
         if zip_file and category and user_prompt:
             # Clear previous images when starting new generation
@@ -168,6 +201,7 @@ def render_bulk_image_generator(uid: str, prefix: str = "ig_img"):
                             None,  # Pass None instead of gallery to avoid appending
                             blur,
                             uid,
+                            private_mode=private_mode,
                         )
                     else:
                         images = process_zip_and_generate_images_multimodel(
@@ -185,6 +219,7 @@ def render_bulk_image_generator(uid: str, prefix: str = "ig_img"):
                             uid,
                             selected_model_key,
                             dynamic_params,
+                            private_mode=private_mode,
                         )
                     st.session_state[gallery_key] = images or []
                     images = st.session_state[gallery_key]
@@ -193,7 +228,25 @@ def render_bulk_image_generator(uid: str, prefix: str = "ig_img"):
                         cols = st.columns(4)
                         for idx, img_path in enumerate(images):
                             with cols[idx % 4]:
-                                st.image(img_path,width='stretch')
+                                if isinstance(img_path, str) and img_path.startswith("data:image"):
+                                    try:
+                                        _, encoded = img_path.split(",", 1)
+                                        img_bytes = base64.b64decode(encoded)
+                                    except Exception:
+                                        img_bytes = None
+                                    if img_bytes:
+                                        st.image(img_bytes, width='stretch')
+                                        st.download_button(
+                                            f"Download image {idx + 1}",
+                                            img_bytes,
+                                            file_name=f"private_image_{idx + 1}.png",
+                                            mime="image/png",
+                                            key=f"{prefix}_private_dl_{idx}"
+                                        )
+                                    else:
+                                        st.warning("Unable to display private image data.")
+                                else:
+                                    st.image(img_path,width='stretch')
                         st.success("Image generation completed!")
                     else:
                         st.info("No images generated.")