Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,3 +1,4 @@
+# main.py
 import os
 import io
 import zipfile
@@ -23,29 +24,38 @@ load_dotenv()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("imagegen_app")
 
+# ----------------------------
+# Config / Constants
+# ----------------------------
 REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
 MONGO_URI = os.getenv("MONGO_URI")
 MONGO_DB = os.getenv("MONGO_DB", "adgenesis_image_text")
 MONGO_COLLECTION = os.getenv("MONGO_COLLECTION", "creatives")
+
 MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
 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"},
-    "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"},
+    "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)
+# ----------------------------
 def get_mongo_collection():
-    if not hasattr(_thread_local, "mongo_collection"):
+    if not hasattr(_thread_local, 'mongo_collection'):
         if not MONGO_URI:
             _thread_local.mongo_collection = None
             return None
@@ -53,7 +63,7 @@ def get_mongo_collection():
             client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000)
             db = client[MONGO_DB]
             collection = db[MONGO_COLLECTION]
-            client.admin.command("ping")
+            client.admin.command('ping')
             _thread_local.mongo_collection = collection
         except Exception as e:
             logger.error(f"MongoDB connection failed: {e}")
@@ -61,8 +71,8 @@ def get_mongo_collection():
     return _thread_local.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"]
+    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:
             _thread_local.s3_client = None
@@ -84,24 +94,31 @@ def get_s3_client():
 def get_model_config(model_key: str) -> Optional[Dict[str, Any]]:
     return MODEL_REGISTRY.get(model_key)
 
+# ----------------------------
+# R2 upload
+# ----------------------------
 def upload_to_r2_optimized(image_bytes: bytes) -> Optional[str]:
     s3_client = get_s3_client()
     if not s3_client:
         return None
     try:
         filename = f"{uuid4().hex}.png"
-        file_key = f"adgenesis_image_file/json/images/{filename}"
+        file_key = f"adgenesis_image_file/balraaj/images/{filename}"
         s3_client.put_object(
             Bucket=os.getenv("R2_BUCKET_NAME"),
             Key=file_key,
             Body=image_bytes,
             ContentType="image/png",
         )
-        return f"{os.getenv('NEW_BASE').rstrip('/')}/{file_key}"
+        r2_url = f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
+        return r2_url
     except Exception as e:
         logger.error(f"S3 upload failed: {e}")
         return None
 
+# ----------------------------
+# Generation & fetching
+# ----------------------------
 def generate_one_image_optimized(model_key: str, prompt: str, aspect_ratio: str) -> List[str]:
     if not REPLICATE_API_TOKEN:
         return []
@@ -113,15 +130,18 @@ def generate_one_image_optimized(model_key: str, prompt: str, aspect_ratio: str)
         ar_param = config["param_name"]
         inputs = {"prompt": prompt, ar_param: aspect_ratio}
         output = replicate.run(model_id, input=inputs)
+        # Normalize to list[str]
         if isinstance(output, list) and output:
-            return [str(output[0])]
+            first = output[0]
+            return [getattr(first, "url", str(first))]
         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 []
+        return []
 
 def fetch_image_bytes_optimized(url: str) -> Optional[bytes]:
     for attempt in range(RETRY_ATTEMPTS):
@@ -135,12 +155,12 @@ def fetch_image_bytes_optimized(url: str) -> Optional[bytes]:
             time.sleep(1)
     return None
 
-def process_single_image(args: Tuple[str,str,str,int]) -> Dict[str,Any]:
+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}
+    result = {"index": index, "success": False, "source_url": None, "r2_url": None, "error": None}
     urls = generate_one_image_optimized(model_key, prompt, aspect_ratio)
     if not urls:
-        result["error"] = "No URLs returned"
+        result["error"] = "No URLs returned from generation"
         return result
     source_url = urls[0]
     result["source_url"] = source_url
@@ -156,136 +176,322 @@ def process_single_image(args: Tuple[str,str,str,int]) -> Dict[str,Any]:
         result["error"] = "Failed to upload to R2"
     return result
 
-def generate_images_parallel(model_key: str, aspect_ratio: str, prompt: str) -> Tuple[List[str],List[str],List[str]]:
-    result = process_single_image((model_key, prompt, aspect_ratio, 0))
-    if result["success"]:
-        return [result["r2_url"]], [result["source_url"]], []
+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 [], [], [result["error"] or "Generation failed"]
+        return [], [], [res["error"] or "Generation failed"]
 
-def save_creative_record_optimized(model_key:str,aspect_ratio:str,prompt:str,urls:List[str]) -> Optional[str]:
+# ----------------------------
+# 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:
         return None
     try:
-        doc = {"model": model_key,"aspect_ratio": aspect_ratio,"prompt": prompt,"urls": urls,"num_images": len(urls),"lob": "json_batch","created_at": datetime.utcnow()}
-        return str(collection.insert_one(doc).inserted_id)
+        doc = {
+            "model": model_key,
+            "aspect_ratio": aspect_ratio,
+            "prompt": prompt,
+            "urls": urls,
+            "num_images": len(urls),
+            "lob": "balraaj",
+            "created_at": datetime.utcnow()
+        }
+        ins = collection.insert_one(doc)
+        return str(ins.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]:
+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
+        return [], 0
     try:
-        total_count = collection.count_documents({"created_at":{"$gte":start_dt,"$lt":end_dt},"lob":"json_batch"})
-        cursor = collection.find({"created_at":{"$gte":start_dt,"$lt":end_dt},"lob":"json_batch"}).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)
+        )
         return list(cursor), total_count
     except Exception:
-        return [],0
+        return [], 0
+
+# ----------------------------
+# UI helpers: images
+# ----------------------------
+@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):
+def display_image_with_download_optimized(url: str):
     try:
-        img_bytes = fetch_image_bytes_optimized(url)
+        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 += ".png"
-        st.download_button("Download image", data=img_bytes, file_name=base, mime="image/png", use_container_width=True)
+            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
-    cols = st.columns(min(4,len(urls)))
-    for i,url in enumerate(urls):
-        with cols[i % len(cols)]:
+    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="images_bundle.zip"):
-    if not urls: return
+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 zipf:
-        for i,url in enumerate(urls,1):
-            img = fetch_image_bytes_optimized(url)
-            if img:
-                base = os.path.basename(urlparse(url).path) or f"img_{i}.png"
-                if not os.path.splitext(base)[1]:
-                    base += ".png"
-                zipf.writestr(base,img)
+    with zipfile.ZipFile(zip_buffer, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
+        for idx, url in enumerate(urls, 1):
+            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)
+    st.download_button(
+        "Download All Images",
+        data=zip_buffer,
+        file_name=filename,
+        mime="application/zip",
+        use_container_width=True
+    )
 
-def load_json_prompts(file)->List[Dict[str,str]]:
-    raw=file.getvalue().decode("utf-8")
-    data=json.loads(raw)
-    if not isinstance(data,dict) or "prompts" not in data or not isinstance(data["prompts"],list):
-        raise ValueError("JSON must be { 'prompts': [ 'string', ... ] }")
-    return [{"id":f"p{i}","content":p.strip()} for i,p in enumerate(data["prompts"],1) if isinstance(p,str) and p.strip()]
+# ----------------------------
+# 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
+
+# ----------------------------
+# JSON page (parallel across prompts)
+# ----------------------------
+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()),0)
-    with col2: default_aspect=st.selectbox("Default Aspect Ratio",MODEL_REGISTRY[default_model]["aspect_ratios"],0)
-    debug=st.checkbox("Debug Mode",False)
+    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=load_json_prompts(up)
-            st.json(prompts)
-            if st.button("Generate for All Prompts",type="primary",use_container_width=True):
-                handle_bulk_json_generation(prompts,default_model,default_aspect,debug)
+            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(str(e))
-
-def handle_bulk_json_generation(prompts:List[Dict[str,str]],default_model:str,default_aspect:str,debug:bool):
-    total=len(prompts)
-    all_urls=[]
-    for i,p in enumerate(prompts,1):
-        st.markdown(f"**Prompt {i}/{total}** — {p['content']}")
-        r2,src,errs=generate_images_parallel(default_model,default_aspect,p["content"])
-        if r2:
-            save_creative_record_optimized(default_model,default_aspect,p["content"],r2)
-            display_image_gallery_optimized(r2)
-            all_urls.extend(r2)
-        elif src:
-            display_image_gallery_optimized(src)
-            all_urls.extend(src)
-        else:
-            st.error("No image generated")
-        if errs and debug:
-            st.error(errs)
+            st.error(f"Failed to read prompts: {e}")
+    else:
+        st.caption('Expected format: { "prompts": ["prompt 1", "prompt 2", ...] }')
+
+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
+    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}")
+
+    # Final all-images gallery & ZIP
     if all_urls:
-        st.subheader("Download All")
-        bulk_download_button(all_urls,"all_prompts_images.zip")
+        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")
 
+# ----------------------------
+# Creative Library page
+# ----------------------------
 def render_library_page():
     st.subheader("Creative Library")
-    today=datetime.utcnow().date()
-    start_date=st.date_input("Start date",today-timedelta(days=30))
-    end_date=st.date_input("End date",today)
-    start_dt=datetime.combine(start_date,datetime.min.time())
-    end_dt=datetime.combine(end_date+timedelta(days=1),datetime.min.time())
-    records,total=query_creatives_optimized(start_dt,end_dt,0)
-    st.caption(f"{total} items")
+    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())
+
+    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
+
+    st.caption(f"Total items: {total_count}")
+    # simple gallery by record
     for rec in records:
-        urls=rec.get("urls",[])
-        if urls: display_image_gallery_optimized(urls)
+        urls = rec.get("urls", []) or []
+        if urls:
+            display_image_gallery_optimized(urls)
+
+# ----------------------------
+# 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."
 
+# ----------------------------
+# App shell
+# ----------------------------
 def main_app():
+    st.set_page_config(page_title="File-to-Image • Creative Library", layout="wide")
     st.title("File-to-Image Generator")
-    page=st.sidebar.radio("Navigation",["Generate from JSON","Creative Library"])
-    if page=="Generate from JSON": render_json_page()
-    else: render_library_page()
+    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()
 
 def main():
-    main_app()
+    if "authenticated" not in st.session_state:
+        st.session_state["authenticated"] = False
+    if not st.session_state["authenticated"]:
+        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)
+    else:
+        main_app()
 
-if __name__=="__main__":
+if __name__ == "__main__":
     main()