Update src/streamlit_app.py

src/streamlit_app.py

@@ -35,7 +35,12 @@ MONGO_COLLECTION = os.getenv("MONGO_COLLECTION", "creatives")
 REQUEST_TIMEOUT = 30
 RETRY_ATTEMPTS = 3
 MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
-LIBRARY_PAGE_SIZE = 200  # aggregate many for one gallery view
+LIBRARY_PAGE_SIZE = 200  # aggregated view
+
+# Global throttle so we don't overload Replicate / network
+# This caps total in-flight image generations across all prompts.
+GLOBAL_CONCURRENCY = max(4, min(16, MAX_WORKERS))
+_GEN_SEMAPHORE = threading.Semaphore(GLOBAL_CONCURRENCY)
 
 MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {
     "imagegen-4-ultra": {"id": "google/imagen-4-ultra","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],"param_name": "aspect_ratio"},
@@ -65,7 +70,6 @@ def show_env_warnings():
 # ----------------------------
 def get_replicate_client():
     if not hasattr(_thread_local, "replicate_client"):
-        # Explicit client avoids env-specific issues with module-level run()
         _thread_local.replicate_client = replicate.Client(api_token=REPLICATE_API_TOKEN) if REPLICATE_API_TOKEN else None
     return _thread_local.replicate_client
 
@@ -160,22 +164,46 @@ def fetch_bytes(url: str) -> Optional[bytes]:
             time.sleep(1)
     return None
 
-def process_prompt(i: int, text: str, model: str, aspect: str) -> Dict[str, Any]:
+def _generate_single_image_full(prompt: str, model: str, aspect: str) -> Tuple[Optional[str], Optional[str]]:
     """
-    One image per prompt:
-    - generate via Replicate
-    - try to upload to R2
-    - fallback to source url if R2 not available
+    Throttled: generate -> fetch -> upload -> return final_url or src_url
+    Returns (final_url_or_src, error_text)
     """
-    urls = generate_one(model, text, aspect)
+    with _GEN_SEMAPHORE:  # throttle
+        urls = generate_one(model, prompt, aspect)
     if not urls:
-        return {"idx": i, "urls": [], "error": "No URLs"}
+        return None, "No URLs"
     src = urls[0]
     data = fetch_bytes(src)
     if data is None:
-        return {"idx": i, "urls": [], "error": "Fetch failed"}
+        return None, "Fetch failed"
     r2 = upload_to_r2(data)
-    return {"idx": i, "urls": [r2 or src], "error": None}
+    return (r2 or src), None
+
+def process_prompt_batch(idx: int, prompt: str, model: str, aspect: str, num_images: int) -> Dict[str, Any]:
+    """
+    Generate N images for one prompt.
+    Uses a small threadpool but still obeys GLOBAL throttle.
+    """
+    n = max(1, int(num_images))
+    urls: List[str] = []
+    errs: List[str] = []
+
+    # limit inner parallelism to keep global pressure sane
+    inner_workers = min( min(4, n), GLOBAL_CONCURRENCY )
+    if n == 1:
+        u, e = _generate_single_image_full(prompt, model, aspect)
+        if u: urls.append(u)
+        if e: errs.append(e)
+    else:
+        with ThreadPoolExecutor(max_workers=inner_workers) as ex:
+            futures = [ex.submit(_generate_single_image_full, prompt, model, aspect) for _ in range(n)]
+            for fut in as_completed(futures):
+                u, e = fut.result()
+                if u: urls.append(u)
+                if e: errs.append(e)
+
+    return {"idx": idx, "urls": urls, "errors": errs}
 
 # ----------------------------
 # Persistence
@@ -190,6 +218,7 @@ def save_record(model: str, aspect: str, prompt: str, urls: List[str]):
             "aspect_ratio": aspect,
             "prompt": prompt,
             "urls": urls,
+            "num_images": len(urls),
             "lob": "balraaj",
             "created_at": datetime.utcnow(),
         }).inserted_id)
@@ -200,7 +229,6 @@ def save_record(model: str, aspect: str, prompt: str, urls: List[str]):
 @st.cache_data(ttl=300)
 def query_records(start: datetime, end: datetime) -> List[Dict[str, Any]]:
     coll = get_mongo_collection()
-    # FIX: compare with None explicitly (avoids NotImplementedError)
     if coll is None:
         return []
     try:
@@ -240,7 +268,7 @@ def bulk_zip(urls: List[str]):
     st.download_button("Download All", buf, "images.zip", "application/zip", use_container_width=True)
 
 # ----------------------------
-# JSON loader & batch run (parallel)
+# JSON loader & batch run (parallel across prompts)
 # ----------------------------
 def load_json(file) -> List[str]:
     data = json.loads(file.getvalue().decode("utf-8"))
@@ -251,7 +279,7 @@ def load_json(file) -> List[str]:
         raise ValueError("No valid prompts found.")
     return out
 
-def run_batch(prompts: List[str], model: str, aspect: str):
+def run_batch(prompts: List[str], model: str, aspect: str, num_images: int):
     total = len(prompts)
     rows = [st.empty() for _ in prompts]
     progress = st.progress(0.0, text=f"0/{total}")
@@ -259,22 +287,25 @@ def run_batch(prompts: List[str], model: str, aspect: str):
     all_urls: List[str] = []
     done = 0
 
-    max_workers = min(MAX_WORKERS, max(2, (os.cpu_count() or 2)))
+    outer_workers = min(MAX_WORKERS, max(2, (os.cpu_count() or 2)))
     with st.spinner("Generating images..."):
-        with ThreadPoolExecutor(max_workers=max_workers) as ex:
-            futs = {ex.submit(process_prompt, i, p, model, aspect): i for i, p in enumerate(prompts, 1)}
+        with ThreadPoolExecutor(max_workers=outer_workers) as ex:
+            futs = {ex.submit(process_prompt_batch, i, p, model, aspect, num_images): i for i, p in enumerate(prompts, 1)}
             for fut in as_completed(futs):
                 i = futs[fut]
                 try:
                     res = fut.result()
                 except Exception as e:
-                    res = {"idx": i, "urls": [], "error": str(e)}
+                    res = {"idx": i, "urls": [], "errors": [str(e)]}
+
                 if res["urls"]:
                     save_record(model, aspect, prompts[i-1], res["urls"])
-                    rows[i-1].success(f"Prompt {i}/{total} ✓")
+                    rows[i-1].success(f"Prompt {i}/{total} ✓  ({len(res['urls'])} images)")
                     all_urls.extend(res["urls"])
                 else:
-                    rows[i-1].error(f"Prompt {i}/{total} ✗ ({res['error']})")
+                    err_msg = ", ".join(res.get("errors") or ["No images"])
+                    rows[i-1].error(f"Prompt {i}/{total} ✗  ({err_msg})")
+
                 done += 1
                 progress.progress(done/total, text=f"{done}/{total}")
 
@@ -291,12 +322,15 @@ def run_batch(prompts: List[str], model: str, aspect: str):
 def render_json_page():
     st.subheader("Generate from JSON")
     show_env_warnings()
+
     up = st.file_uploader("Upload JSON", type=["json"])
-    col1, col2 = st.columns([1, 1])
-    with col1:
+    c1, c2, c3 = st.columns([1, 1, 1])
+    with c1:
         model = st.selectbox("Model", list(MODEL_REGISTRY.keys()), 0)
-    with col2:
+    with c2:
         aspect = st.selectbox("Aspect", MODEL_REGISTRY[model]["aspect_ratios"], 0)
+    with c3:
+        num_images = st.slider("Images per prompt", min_value=1, max_value=12, value=1, step=1)
 
     if up:
         try:
@@ -304,7 +338,7 @@ def render_json_page():
             with st.expander("Preview prompts", expanded=False):
                 st.json(prompts)
             if st.button("Generate", type="primary", use_container_width=True):
-                run_batch(prompts, model, aspect)
+                run_batch(prompts, model, aspect, num_images)
         except Exception as e:
             st.error(str(e))
     else:
@@ -339,7 +373,7 @@ def render_library_page():
 def check_token(tok: str) -> bool:
     acc = os.getenv("ACCESS_TOKEN")
     if not acc:
-        # If ACCESS_TOKEN is not configured, allow through to avoid lockout in dev.
+        # No token configured → allow dev access
         return True
     return tok == acc