Update app_pages/image_generation.py

app_pages/image_generation.py

@@ -28,7 +28,8 @@ def _zip_gallery_images(gallery_items):
                 zipf.write(os.path.join(temp_dir, file_name), arcname=file_name)
         return zip_path
     except Exception as e:
-        logger.critical(f"Failed to create zip: {e}"); return None
+        logger.critical(f"Failed to create zip: {e}")
+        return None
     finally:
         shutil.rmtree(temp_dir, ignore_errors=True)
 
@@ -54,31 +55,64 @@ def render_bulk_image_generator(prefix: str = "ig_img"):
 
     if demo_btn or gen_all_btn:
         if zip_file and category and user_prompt:
-            temp_dir = tempfile.mkdtemp()
-            try:
-                temp_path = os.path.join(temp_dir, zip_file.name if zip_file.name.lower().endswith(".zip") else f"input.{zip_file.name.split('.')[-1]}")
-                with open(temp_path, "wb") as f: f.write(zip_file.read())
-                images = process_zip_and_generate_images(temp_path, category, size, quality, user_prompt, sentiment, platform, num_images, demo_btn, gallery, blur)
-                st.session_state[gallery_key] = images or []; images = st.session_state[gallery_key]
-                if images:
-                    cols = st.columns(4)
-                    for idx, img_path in enumerate(images):
-                        with cols[idx % 4]: st.image(img_path, use_container_width=True)
-                else:
-                    st.info("No images generated.")
-            except Exception as e:
-                st.error(f"Error: {e}"); logger.exception("Generation failed (embedded).")
-            finally:
-                shutil.rmtree(temp_dir, ignore_errors=True)
+            with st.spinner(" Generating image variations..."):
+                temp_dir = tempfile.mkdtemp()
+                try:
+                    temp_path = os.path.join(
+                        temp_dir,
+                        zip_file.name if zip_file.name.lower().endswith(".zip")
+                        else f"input.{zip_file.name.split('.')[-1]}"
+                    )
+                    with open(temp_path, "wb") as f:
+                        f.write(zip_file.read())
+
+                    images = process_zip_and_generate_images(
+                        temp_path,
+                        category,
+                        size,
+                        quality,
+                        user_prompt,
+                        sentiment,
+                        platform,
+                        num_images,
+                        demo_btn,
+                        gallery,
+                        blur,
+                    )
+                    st.session_state[gallery_key] = images or []
+                    images = st.session_state[gallery_key]
+
+                    if images:
+                        cols = st.columns(4)
+                        for idx, img_path in enumerate(images):
+                            with cols[idx % 4]:
+                                st.image(img_path, use_container_width=True)
+                        st.success("Image generation completed!")
+                    else:
+                        st.info("No images generated.")
+                except Exception as e:
+                    st.error(f"Error: {e}")
+                    logger.exception("Generation failed (embedded).")
+                finally:
+                    shutil.rmtree(temp_dir, ignore_errors=True)
         else:
             st.warning("Please upload a file and fill all required fields.")
 
     if download_btn:
-        zip_path = _zip_gallery_images(st.session_state.get(gallery_key, []))
-        if zip_path:
-            with open(zip_path, "rb") as f:
-                st.download_button("Download ZIP", data=f.read(), file_name="generated_images.zip", mime="application/zip", key=f"{prefix}_zip_dl_btn", use_container_width=True)
-            try: os.remove(zip_path)
-            except Exception: pass
-        else:
-            st.warning("No images to zip and download.")
+            zip_path = _zip_gallery_images(st.session_state.get(gallery_key, []))
+            if zip_path:
+                with open(zip_path, "rb") as f:
+                    st.download_button(
+                        "Download ZIP",
+                        data=f.read(),
+                        file_name="generated_images.zip",
+                        mime="application/zip",
+                        key=f"{prefix}_zip_dl_btn",
+                        use_container_width=True
+                    )
+                try:
+                    os.remove(zip_path)
+                except Exception:
+                    pass
+            else:
+                st.warning("No images to zip and download.")