Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,7 +7,7 @@ import time
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta, date
-from typing import Dict, Any, List, Tuple, Optional, Union
+from typing import Dict, Any, List, Tuple, Optional
 import requests
 import streamlit as st
 from pymongo import MongoClient
@@ -37,22 +37,18 @@ REQUEST_TIMEOUT = 30
 RETRY_ATTEMPTS = 3
 LIBRARY_PAGE_SIZE = 20
 
-# Model registry (subset with common ARs)
 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"},
     "imagen-4":          {"id": "google/imagen-4",       "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"], "param_name": "aspect_ratio"},
-    "nano-banana":       {"id": "google/nano-banana",    "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"], "param_name": "aspect_ratio"},
     "qwen":              {"id": "qwen/qwen-image",       "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3"], "param_name": "aspect_ratio"},
     "seedream-3":        {"id": "bytedance/seedream-3",  "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3","21:9"], "param_name": "aspect_ratio"},
     "recraft-v3":        {"id": "recraft-ai/recraft-v3", "aspect_ratios": ["1:1","4:3","3:4","3:2","2:3","16:9","9:16"], "param_name": "aspect_ratio"},
-    "photon":            {"id": "luma/photon",           "aspect_ratios": ["1:1","3:4","4:3","9:16","16:9","21:9"], "param_name": "aspect_ratio"},
-    "ideogram-v3-quality":{"id": "ideogram-ai/ideogram-v3-quality", "aspect_ratios": ["1:1","16:9","9:16","2:3","3:2","4:5","5:4"], "param_name": "aspect_ratio"},
 }
 
 _thread_local = threading.local()
 
 # ----------------------------
-# Infra helpers (Mongo / S3)
+# Infra helpers
 # ----------------------------
 def get_mongo_collection():
     if not hasattr(_thread_local, 'mongo_collection'):
@@ -73,8 +69,7 @@ def get_mongo_collection():
 def get_s3_client():
     if not hasattr(_thread_local, 's3_client'):
         required_vars = ["R2_ENDPOINT", "R2_ACCESS_KEY", "R2_SECRET_KEY", "R2_BUCKET_NAME"]
-        missing = [var for var in required_vars if not os.getenv(var)]
-        if missing:
+        if any(not os.getenv(v) for v in required_vars):
             _thread_local.s3_client = None
             return None
         try:
@@ -95,403 +90,213 @@ def get_model_config(model_key: str) -> Optional[Dict[str, Any]]:
     return MODEL_REGISTRY.get(model_key)
 
 # ----------------------------
-# R2 upload
+# Upload & generation
 # ----------------------------
-def upload_to_r2_optimized(image_bytes: bytes) -> Optional[str]:
-    s3_client = get_s3_client()
-    if not s3_client:
+def upload_to_r2(image_bytes: bytes) -> Optional[str]:
+    s3 = get_s3_client()
+    if not s3:
         return None
     try:
         filename = f"{uuid4().hex}.png"
         file_key = f"adgenesis_image_file/balraaj/images/{filename}"
-        s3_client.put_object(
+        s3.put_object(
             Bucket=os.getenv("R2_BUCKET_NAME"),
             Key=file_key,
             Body=image_bytes,
             ContentType="image/png",
         )
-        r2_url = f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
-        return r2_url
+        return f"{os.getenv('NEW_BASE').rstrip('/')}/{file_key}"
     except Exception as e:
-        logger.error(f"S3 upload failed: {e}")
+        logger.error(f"Upload failed: {e}")
         return None
 
-# ----------------------------
-# Generation & fetching
-# ----------------------------
-def generate_one_image_optimized(model_key: str, prompt: str, aspect_ratio: str) -> List[str]:
+def generate_one(model_key: str, prompt: str, aspect_ratio: str) -> List[str]:
     if not REPLICATE_API_TOKEN:
         return []
     config = get_model_config(model_key)
     if not config:
         return []
     try:
-        model_id = config["id"]
-        ar_param = config["param_name"]
-        inputs = {"prompt": prompt, ar_param: aspect_ratio}
-        output = replicate.run(model_id, input=inputs)
-        # Normalize to list[str]
+        output = replicate.run(config["id"], input={"prompt": prompt, config["param_name"]: aspect_ratio})
         if isinstance(output, list) and output:
-            first = output[0]
-            return [getattr(first, "url", str(first))]
+            return [str(output[0])]
         elif isinstance(output, str):
             return [output]
-        elif hasattr(output, "url"):
-            return [getattr(output, "url")]
         return []
     except Exception as e:
         logger.error(f"Replicate error: {e}")
         return []
 
-def fetch_image_bytes_optimized(url: str) -> Optional[bytes]:
-    for attempt in range(RETRY_ATTEMPTS):
+def fetch_bytes(url: str) -> Optional[bytes]:
+    for _ in range(RETRY_ATTEMPTS):
         try:
-            response = requests.get(url, timeout=REQUEST_TIMEOUT, stream=True)
-            response.raise_for_status()
-            return response.content
+            r = requests.get(url, timeout=REQUEST_TIMEOUT, stream=True)
+            r.raise_for_status()
+            return r.content
         except Exception:
-            if attempt == RETRY_ATTEMPTS - 1:
-                return None
             time.sleep(1)
     return None
 
-def process_single_image(args: Tuple[str, str, str, int]) -> Dict[str, Any]:
-    model_key, prompt, aspect_ratio, index = args
-    result = {"index": index, "success": False, "source_url": None, "r2_url": None, "error": None}
-    urls = generate_one_image_optimized(model_key, prompt, aspect_ratio)
+def process_prompt(i: int, text: str, model: str, aspect: str):
+    urls = generate_one(model, text, aspect)
     if not urls:
-        result["error"] = "No URLs returned from generation"
-        return result
-    source_url = urls[0]
-    result["source_url"] = source_url
-    img_bytes = fetch_image_bytes_optimized(source_url)
+        return {"idx": i, "urls": [], "error": "No URLs"}
+    img_bytes = fetch_bytes(urls[0])
     if not img_bytes:
-        result["error"] = "Failed to fetch image bytes"
-        return result
-    r2_url = upload_to_r2_optimized(img_bytes)
-    if r2_url:
-        result["r2_url"] = r2_url
-        result["success"] = True
-    else:
-        result["error"] = "Failed to upload to R2"
-    return result
-
-def generate_one_per_prompt(model_key: str, aspect_ratio: str, prompt: str) -> Tuple[List[str], List[str], List[str]]:
-    """One image per prompt (no parallel within a prompt)."""
-    res = process_single_image((model_key, prompt, aspect_ratio, 0))
-    if res["success"]:
-        return [res["r2_url"]], [res["source_url"]], []
-    else:
-        return [], [], [res["error"] or "Generation failed"]
+        return {"idx": i, "urls": [], "error": "Fetch failed"}
+    r2 = upload_to_r2(img_bytes)
+    return {"idx": i, "urls": [r2 or urls[0]], "error": None}
 
 # ----------------------------
 # Persistence
 # ----------------------------
-def save_creative_record_optimized(model_key: str, aspect_ratio: str, prompt: str, urls: List[str]) -> Optional[str]:
-    collection = get_mongo_collection()
-    if collection is None:
+def save_record(model: str, aspect: str, prompt: str, urls: List[str]):
+    coll = get_mongo_collection()
+    if not coll:
         return None
     try:
-        doc = {
-            "model": model_key,
-            "aspect_ratio": aspect_ratio,
+        return str(coll.insert_one({
+            "model": model,
+            "aspect_ratio": aspect,
             "prompt": prompt,
             "urls": urls,
-            "num_images": len(urls),
             "lob": "balraaj",
             "created_at": datetime.utcnow()
-        }
-        ins = collection.insert_one(doc)
-        return str(ins.inserted_id)
+        }).inserted_id)
     except Exception as e:
         logger.error(f"Mongo insert failed: {e}")
         return None
 
 @st.cache_data(ttl=300)
-def query_creatives_optimized(start_dt: datetime, end_dt: datetime, page: int = 0) -> Tuple[List[Dict[str, Any]], int]:
-    collection = get_mongo_collection()
-    if collection is None:
-        return [], 0
+def query_records(start: datetime, end: datetime) -> List[Dict[str, Any]]:
+    coll = get_mongo_collection()
+    if not coll:
+        return []
     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)
-        )
-        return list(cursor), total_count
+        return list(coll.find(
+            {"created_at": {"$gte": start, "$lt": end}, "lob": "balraaj"}
+        ).sort("created_at", -1).limit(LIBRARY_PAGE_SIZE))
     except Exception:
-        return [], 0
+        return []
 
 # ----------------------------
-# UI helpers: images
+# Gallery helpers
 # ----------------------------
-@st.cache_data(ttl=3600)
-def get_image_bytes_cached(url: str) -> Optional[bytes]:
-    return fetch_image_bytes_optimized(url)
-
-def display_image_with_download_optimized(url: str):
-    try:
-        img_bytes = get_image_bytes_cached(url)
-        if not img_bytes:
-            st.error("Failed to load image")
-            return
-        st.image(img_bytes, use_container_width=True)
-        base = os.path.basename(urlparse(url).path) or "image.png"
-        if not os.path.splitext(base)[1]:
-            base = f"{base}.png"
-        st.download_button(
-            label="Download image",
-            data=img_bytes,
-            file_name=base,
-            mime="image/png",
-            use_container_width=True
-        )
-    except Exception as e:
-        st.error(f"Failed to display image: {e}")
-
-def display_image_gallery_optimized(urls: List[str]):
-    if not urls: 
-        return
-    num_cols = min(4, max(1, len(urls)))
-    cols = st.columns(num_cols)
-    for i, url in enumerate(urls):
-        with cols[i % num_cols]:
-            display_image_with_download_optimized(url)
-
-def bulk_download_button(urls: List[str], filename: str = "images_bundle.zip"):
-    if not urls:
-        return
-    zip_buffer = io.BytesIO()
-    with zipfile.ZipFile(zip_buffer, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
-        for idx, url in enumerate(urls, 1):
+def display_gallery(urls: List[str]):
+    if not urls: return
+    cols = st.columns(4)
+    for i, u in enumerate(urls):
+        with cols[i % 4]:
             try:
-                img_bytes = fetch_image_bytes_optimized(url)
-                if img_bytes:
-                    path = urlparse(url).path
-                    base = os.path.basename(path) or f"image_{idx}.png"
-                    if not os.path.splitext(base)[1]:
-                        base = f"{base}.png"
-                    zip_file.writestr(base, img_bytes)
-            except Exception:
-                pass
-    zip_buffer.seek(0)
-    st.download_button(
-        "Download All Images",
-        data=zip_buffer,
-        file_name=filename,
-        mime="application/zip",
-        use_container_width=True
-    )
-
-# ----------------------------
-# JSON loader (STRICT)
-# ----------------------------
-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]] = []
-    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
+                img = fetch_bytes(u)
+                if img:
+                    st.image(img, use_container_width=True)
+            except:
+                st.error("Failed")
+
+def bulk_zip(urls: List[str]):
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as z:
+        for i, u in enumerate(urls, 1):
+            data = fetch_bytes(u)
+            if data:
+                name = f"image_{i}.png"
+                z.writestr(name, data)
+    buf.seek(0)
+    st.download_button("Download All", buf, "images.zip", "application/zip")
 
 # ----------------------------
-# JSON page (parallel across prompts)
+# JSON loader & runner
 # ----------------------------
-def _run_single_prompt(idx: int, prompt_text: str, model_key: str, aspect_ratio: str):
-    r2_urls, src_urls, gen_errors = generate_one_per_prompt(model_key, aspect_ratio, prompt_text)
-    rec_id = None
-    if r2_urls:
-        rec_id = save_creative_record_optimized(model_key, aspect_ratio, prompt_text, r2_urls)
-    return {
-        "idx": idx,
-        "prompt": prompt_text,
-        "r2_urls": r2_urls,
-        "src_urls": src_urls,
-        "errors": gen_errors,
-        "rec_id": rec_id,
-    }
-
-def render_json_page():
-    st.subheader("Generate from JSON Prompts")
-    up = st.file_uploader("Upload prompts JSON", type=["json"])
-
-    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")
-
-    debug_mode = st.checkbox("Debug Mode", value=False, key="json_debug")
-
-    if up:
-        try:
-            prompts_list = load_json_prompts(up)
-            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_parallel(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('Expected format: { "prompts": ["prompt 1", "prompt 2", ...] }')
+def load_json(file) -> List[str]:
+    data = json.loads(file.getvalue().decode("utf-8"))
+    if not isinstance(data, dict) or "prompts" not in data:
+        raise ValueError("JSON must be { 'prompts': [ ... ] }")
+    return [p for p in data["prompts"] if isinstance(p, str) and p.strip()]
 
-def handle_bulk_json_generation_parallel(prompts: List[Dict[str, str]], default_model: str, default_aspect: str, debug: bool):
-    if not REPLICATE_API_TOKEN:
-        st.error("Missing REPLICATE_API_TOKEN. Set it as an environment variable.")
-        return
+def run_batch(prompts: List[str], model: str, aspect: str):
     total = len(prompts)
-    if total == 0:
-        st.info("No prompts to process.")
-        return
-
-    # Placeholders for stable on-page order
-    blocks = [st.container(border=True) for _ in range(total)]
-    progress = st.progress(0, text=f"Starting batch • 0/{total}")
-
-    all_urls: List[str] = []
-    completed = 0
-
-    max_workers = min(MAX_WORKERS, max(2, (os.cpu_count() or 2)))
-
-    with st.spinner("Generating images..."):
-        futures = {}
-        with ThreadPoolExecutor(max_workers=max_workers) as ex:
-            for i, p in enumerate(prompts, 1):
-                prompt_text = p.get("content", "").strip()
-                if not prompt_text:
-                    # render immediately as invalid
-                    with blocks[i-1]:
-                        st.markdown(f"**Prompt {i}/{total}** — (empty)")
-                        st.error("Prompt text is empty. Skipping.")
-                    completed += 1
-                    progress.progress(completed / total, text=f"Processed {completed}/{total}")
-                    continue
-                futures[ex.submit(_run_single_prompt, i, prompt_text, default_model, default_aspect)] = i
-
-            for fut in as_completed(futures):
-                i = futures[fut]
-                try:
-                    res = fut.result()
-                except Exception as e:
-                    res = {"idx": i, "prompt": "", "r2_urls": [], "src_urls": [], "errors": [str(e)], "rec_id": None}
-
-                with blocks[i-1]:
-                    st.markdown(f"**Prompt {i}/{total}** — Model: `{default_model}` • Aspect: `{default_aspect}` • Num: `1`")
-                    st.code(res.get("prompt") or "(empty)", language="markdown")
-
-                    if res["r2_urls"]:
-                        st.success(f"Generated 1 image. DB: {res['rec_id'] or 'N/A'}")
-                        display_image_gallery_optimized(res["r2_urls"])
-                        bulk_download_button(res["r2_urls"], filename=f"prompt_{i}_image.zip")
-                        all_urls.extend(res["r2_urls"])
-                    elif res["src_urls"]:
-                        st.warning("Image generated but R2 upload failed. Showing original:")
-                        display_image_gallery_optimized(res["src_urls"])
-                        bulk_download_button(res["src_urls"], filename=f"prompt_{i}_image.zip")
-                        all_urls.extend(res["src_urls"])
-                    else:
-                        st.error("No image was generated for this prompt.")
-
-                    if res.get("errors") and debug:
-                        for e in res["errors"]:
-                            st.error(e)
-
-                completed += 1
-                progress.progress(completed / total, text=f"Processed {completed}/{total}")
+    status = [st.empty() for _ in prompts]
+    progress = st.progress(0, f"0/{total}")
+
+    all_urls = []
+    with ThreadPoolExecutor(max_workers=min(MAX_WORKERS, total)) as ex:
+        futs = {ex.submit(process_prompt, i, p, model, aspect): i for i,p in enumerate(prompts,1)}
+        done = 0
+        for f in as_completed(futs):
+            i = futs[f]
+            try:
+                res = f.result()
+            except Exception as e:
+                res = {"idx": i, "urls": [], "error": str(e)}
+            if res["urls"]:
+                save_record(model, aspect, prompts[i-1], res["urls"])
+                status[i-1].success(f"Prompt {i}/{total} ✓")
+                all_urls.extend(res["urls"])
+            else:
+                status[i-1].error(f"Prompt {i}/{total} ✗ ({res['error']})")
+            done += 1
+            progress.progress(done/total, f"{done}/{total}")
 
-    # Final all-images gallery & ZIP
     if all_urls:
-        st.subheader("All Images Gallery")
-        display_image_gallery_optimized(all_urls)
-        st.subheader("Download All Generated")
-        bulk_download_button(all_urls, filename="all_prompts_images.zip")
+        st.subheader("Gallery")
+        display_gallery(all_urls)
+        bulk_zip(all_urls)
 
 # ----------------------------
-# Creative Library page
+# Pages
 # ----------------------------
-def render_library_page():
-    st.subheader("Creative Library")
-    if "library_page" not in st.session_state:
-        st.session_state.library_page = 0
-
-    today_utc = datetime.utcnow().date()
-    default_start = today_utc - timedelta(days=30)
-
-    c1, c2, c3 = st.columns([1, 1, 1])
-    with c1:
-        start_date: date = st.date_input("Start date", value=default_start)
-    with c2:
-        end_date: date = st.date_input("End date", value=today_utc)
-    with c3:
-        if st.button("Apply Filters", use_container_width=True):
-            st.session_state.library_page = 0
-            st.cache_data.clear()
-
-    start_dt = datetime.combine(start_date, datetime.min.time())
-    end_dt = datetime.combine(end_date + timedelta(days=1), datetime.min.time())
+def render_json_page():
+    st.subheader("Generate from JSON")
+    up = st.file_uploader("Upload JSON", type=["json"])
+    col1,col2 = st.columns([1,1])
+    with col1: model = st.selectbox("Model", list(MODEL_REGISTRY.keys()), 0)
+    with col2: aspect = st.selectbox("Aspect", MODEL_REGISTRY[model]["aspect_ratios"], 0)
 
-    records, total_count = query_creatives_optimized(start_dt, end_dt, st.session_state.library_page)
-    if not records and st.session_state.library_page == 0:
-        st.info("No creatives found for the selected dates.")
-        return
+    if up:
+        prompts = load_json(up)
+        st.json(prompts)
+        if st.button("Generate", type="primary", use_container_width=True):
+            run_batch(prompts, model, aspect)
 
-    st.caption(f"Total items: {total_count}")
-    # simple gallery by record
-    for rec in records:
-        urls = rec.get("urls", []) or []
-        if urls:
-            display_image_gallery_optimized(urls)
+def render_library_page():
+    st.subheader("Creative Library")
+    today = datetime.utcnow().date()
+    start = st.date_input("Start", today - timedelta(days=30))
+    end = st.date_input("End", today)
+    records = query_records(datetime.combine(start, datetime.min.time()),
+                            datetime.combine(end+timedelta(days=1), datetime.min.time()))
+    all_urls = []
+    for r in records:
+        all_urls.extend(r.get("urls", []))
+    if all_urls:
+        display_gallery(all_urls)
+        bulk_zip(all_urls)
+    else:
+        st.info("No records found.")
 
 # ----------------------------
 # Auth
 # ----------------------------
 @lru_cache(maxsize=1)
-def check_token_cached(user_token: str) -> Tuple[bool, str]:
-    ACCESS_TOKEN = os.getenv("ACCESS_TOKEN")
-    if not ACCESS_TOKEN:
-        return False, "Server error: Access token not configured."
-    if user_token == ACCESS_TOKEN:
-        return True, ""
-    return False, "Invalid token."
+def check_token(tok: str):
+    return tok == os.getenv("ACCESS_TOKEN")
 
-# ----------------------------
-# App shell
-# ----------------------------
 def main_app():
-    st.set_page_config(page_title="File-to-Image • Creative Library", layout="wide")
     st.title("File-to-Image Generator")
-    with st.sidebar:
-        page = st.radio("Navigation", ["Generate from JSON", "Creative Library"], index=0)
-    if page == "Generate from JSON":
-        render_json_page()
-    else:
-        render_library_page()
+    page = st.sidebar.radio("Menu", ["Generate from JSON","Creative Library"])
+    if page=="Generate from JSON": render_json_page()
+    else: render_library_page()
 
 def main():
-    if "authenticated" not in st.session_state:
-        st.session_state["authenticated"] = False
-    if not st.session_state["authenticated"]:
+    if not st.session_state.get("auth"):
         st.markdown("## Access Required")
-        token_input = st.text_input("Enter Access Token", type="password")
-        if st.button("Unlock App"):
-            ok, error_msg = check_token_cached(token_input)
-            if ok:
-                st.session_state["authenticated"] = True
-                st.rerun()
-            else:
-                st.error(error_msg)
+        t = st.text_input("Token", type="password")
+        if st.button("Unlock"):
+            if check_token(t): st.session_state["auth"]=True; st.rerun()
+            else: st.error("Invalid token")
     else:
         main_app()
 
-if __name__ == "__main__":
+if __name__=="__main__":
     main()