Update src/streamlit_app.py

src/streamlit_app.py

@@ -213,7 +213,15 @@ def save_creative_record_optimized(model_key: str, aspect_ratio: str, prompt: st
     if collection is None:
         return None
     try:
-        doc = {"model": model_key,"aspect_ratio": aspect_ratio,"prompt": prompt,"urls": urls,"num_images": len(urls),"lob": "balraaj","created_at": datetime.utcnow()}
+        doc = {
+            "model": model_key,
+            "aspect_ratio": aspect_ratio,
+            "prompt": prompt,
+            "urls": urls,
+            "num_images": len(urls),
+            "lob": "balraaj",
+            "created_at": datetime.utcnow()
+        }
         result = collection.insert_one(doc)
         return str(result.inserted_id)
     except Exception:
@@ -225,8 +233,13 @@ def query_creatives_optimized(start_dt: datetime, end_dt: datetime, page: int =
     if collection is None:
         return [], 0
     try:
-        total_count = collection.count_documents({"created_at": {"$gte": start_dt, "$lt": end_dt},"lob": "balraaj"})
-        cursor = collection.find({"created_at": {"$gte": start_dt, "$lt": end_dt}, "lob": "balraaj"}).sort("created_at", -1).skip(page * LIBRARY_PAGE_SIZE).limit(LIBRARY_PAGE_SIZE)
+        total_count = collection.count_documents({"created_at": {"$gte": start_dt, "$lt": end_dt}, "lob": "balraaj"})
+        cursor = (
+            collection.find({"created_at": {"$gte": start_dt, "$lt": end_dt}, "lob": "balraaj"})
+            .sort("created_at", -1)
+            .skip(page * LIBRARY_PAGE_SIZE)
+            .limit(LIBRARY_PAGE_SIZE)
+        )
         records = list(cursor)
         return records, total_count
     except Exception:
@@ -396,97 +409,103 @@ def check_token_cached(user_token: str) -> Tuple[bool, str]:
 def load_json_prompts(file) -> List[Dict[str, Any]]:
     raw = file.getvalue().decode("utf-8", errors="replace")
     data = json.loads(raw)
+
+    if not isinstance(data, dict) or "prompts" not in data or not isinstance(data["prompts"], list):
+        raise ValueError("Invalid JSON. Expected an object with a 'prompts' array of strings.")
+
     prompts_out: List[Dict[str, Any]] = []
-    def push(idx: int, item: Any):
-        if isinstance(item, str):
-            prompts_out.append({"id": f"p{idx}", "content": item})
-        elif isinstance(item, dict) and "content" in item:
-            obj = {"id": str(item.get("id", f"p{idx}")), "content": str(item["content"])}
-            if "num" in item: obj["num"] = int(item["num"])
-            if "aspect_ratio" in item: obj["aspect_ratio"] = str(item["aspect_ratio"])
-            if "model" in item: obj["model"] = str(item["model"])
-            prompts_out.append(obj)
-    if isinstance(data, dict) and "prompts" in data and isinstance(data["prompts"], list):
-        for i, item in enumerate(data["prompts"], 1):
-            push(i, item)
-    elif isinstance(data, list):
-        for i, item in enumerate(data, 1):
-            push(i, item)
+    for i, item in enumerate(data["prompts"], 1):
+        if not isinstance(item, str) or not item.strip():
+            raise ValueError(f"'prompts[{i-1}]' must be a non-empty string.")
+        prompts_out.append({"id": f"p{i}", "content": item.strip()})
+
     return prompts_out
 
 def render_json_page():
     st.subheader("Generate from JSON Prompts")
     up = st.file_uploader("Upload prompts JSON", type=["json"])
-    col1, col2, col3 = st.columns([1,1,1])
+
+    col1, col2 = st.columns([1, 1])
     with col1:
         default_model = st.selectbox("Default Model", list(MODEL_REGISTRY.keys()), index=0)
     with col2:
         aspect_options = MODEL_REGISTRY[default_model]["aspect_ratios"]
         default_aspect = st.selectbox("Default Aspect Ratio", aspect_options, index=0, key="json_default_ar")
-    with col3:
-        default_num = st.slider("Default Images per Prompt", 1, 20, 1, 1, key="json_default_num")
+
     debug_mode = st.checkbox("Debug Mode", value=False, key="json_debug")
+
     if up:
         try:
             prompts_list = load_json_prompts(up)
             if not prompts_list:
                 st.error("No prompts found in the JSON.")
                 return
+
             with st.expander("Preview normalized prompts", expanded=False):
                 st.json(prompts_list, expanded=False)
+
             if st.button("Generate for All Prompts", type="primary", use_container_width=True):
-                handle_bulk_json_generation(prompts_list, default_model, default_aspect, default_num, debug_mode)
+                handle_bulk_json_generation(prompts_list, default_model, default_aspect, debug_mode)
         except json.JSONDecodeError as e:
             st.error(f"Invalid JSON: {e}")
         except Exception as e:
             st.error(f"Failed to read prompts: {e}")
     else:
-        st.caption("Your JSON can be a list of strings, a list of objects with `content`, or `{ 'prompts': [...] }`. Optional per-prompt keys: `model`, `aspect_ratio`, `num`.")
+        st.caption('Expected format: { "prompts": ["prompt 1", "prompt 2", ...] }')
 
-def handle_bulk_json_generation(prompts: List[Dict[str, Any]], default_model: str, default_aspect: str, default_num: int, debug_mode: bool):
+def handle_bulk_json_generation(prompts: List[Dict[str, Any]], default_model: str, default_aspect: str, debug_mode: bool):
     if not REPLICATE_API_TOKEN:
         st.error("Missing REPLICATE_API_TOKEN. Set it as an environment variable.")
         return
+
     total = len(prompts)
     overall_progress = st.progress(0, text=f"Starting batch • 0/{total}")
     all_generated_urls: List[str] = []
     errors_total: List[str] = []
     start_time = time.time()
+
     for idx, p in enumerate(prompts, 1):
-        model_key = str(p.get("model", default_model))
-        aspect_ratio = str(p.get("aspect_ratio", default_aspect))
-        num_images = int(p.get("num", default_num))
+        model_key = default_model
+        aspect_ratio = default_aspect
+        num_images = 1  # exactly one image per prompt
         prompt_text = str(p.get("content", "")).strip()
+
         block = st.container(border=True)
         with block:
-            st.markdown(f"**Prompt {idx}/{total}** — Model: `{model_key}` • Aspect: `{aspect_ratio}` • Num: `{num_images}`")
+            st.markdown(f"**Prompt {idx}/{total}** — Model: `{model_key}` • Aspect: `{aspect_ratio}` • Num: `1`")
             st.code(prompt_text or "(empty)", language="markdown")
+
             if not prompt_text:
                 st.error("Prompt text is empty. Skipping.")
                 overall_progress.progress(min(idx / total, 1.0), text=f"Processed {idx}/{total}")
                 continue
+
             r2_urls, src_urls, gen_errors = generate_images_parallel(model_key, aspect_ratio, prompt_text, num_images)
             rec_id = None
             if r2_urls:
                 rec_id = save_creative_record_optimized(model_key, aspect_ratio, prompt_text, r2_urls)
+
             if r2_urls:
-                st.success(f"Generated {len(r2_urls)} image(s). DB: {rec_id or 'N/A'}")
+                st.success(f"Generated 1 image. DB: {rec_id or 'N/A'}")
                 display_image_gallery_optimized(r2_urls)
-                bulk_download_button(r2_urls, filename=f"prompt_{idx}_images.zip")
+                bulk_download_button(r2_urls, filename=f"prompt_{idx}_image.zip")
                 all_generated_urls.extend(r2_urls)
             elif src_urls:
-                st.warning("Images generated but R2 upload failed. Showing originals:")
+                st.warning("Image generated but R2 upload failed. Showing original:")
                 display_image_gallery_optimized(src_urls)
-                bulk_download_button(src_urls, filename=f"prompt_{idx}_images.zip")
+                bulk_download_button(src_urls, filename=f"prompt_{idx}_image.zip")
                 all_generated_urls.extend(src_urls)
             else:
-                st.error("No images were generated for this prompt.")
+                st.error("No image was generated for this prompt.")
+
             if gen_errors and debug_mode:
                 with st.expander("Errors", expanded=False):
                     for e in gen_errors:
                         st.error(e)
                     errors_total.extend(gen_errors)
+
         overall_progress.progress(min(idx / total, 1.0), text=f"Processed {idx}/{total}")
+
     elapsed = time.time() - start_time
     st.success(f"Batch complete in {elapsed:.1f}s. Total prompts: {total}.")
     if all_generated_urls: