Update app.py

app.py

@@ -1,296 +1,192 @@
 import os
-import math
-import random
 import tempfile
-import gradio as gr
+from functools import lru_cache
 from huggingface_hub import HfApi, hf_hub_download
-
-# ======================================================
-# CONFIG
-# ======================================================
-DEFAULT_REPO = "rahul7star/WAN22-Feb-2026"
-FREE_REPO = "rahul7star/ohamlab"
-FREE_IMAGE_DIR = "showcase/image"
+import gradio as gr
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-APP_PASSWORD = os.getenv("APP_PASSWORD")
-
-if not HF_TOKEN:
-    raise RuntimeError("HF_TOKEN not set")
-if not APP_PASSWORD:
-    raise RuntimeError("APP_PASSWORD not set")
-
-api = HfApi()
-PAGE_SIZE = 6
-
-tmp_dir = tempfile.mkdtemp()
-images_cache = {}     # (repo, date) -> list[(img, prompt)]
-prompt_cache = {}     # folder_path -> prompt text
-
-# ======================================================
-# PROMPT LOADER (NEW)
-# ======================================================
-def load_prompt_from_folder(repo, folder):
-    """
-    Tries to find prompt text inside the same upload folder.
-    Cached for performance.
-    """
-    key = (repo, folder)
-    if key in prompt_cache:
-        return prompt_cache[key]
-
-    files = api.list_repo_files(repo, token=HF_TOKEN)
-    prompt = ""
-
-    for f in files:
-        if not f.startswith(folder + "/"):
-            continue
-
-        name = os.path.basename(f).lower()
 
-        if name in ("prompt.txt", "caption.txt"):
-            try:
-                path = hf_hub_download(repo, f, token=HF_TOKEN)
-                with open(path, "r", encoding="utf-8", errors="ignore") as fp:
-                    prompt = fp.read().strip()
-            except Exception as e:
-                print("[PROMPT ERROR]", e)
-            break
+DEFAULT_REPO = "rahul7star/WAN22-Feb-2026"
+IMAGES_PER_PAGE = 6
 
-    prompt_cache[key] = prompt
-    return prompt
+api = HfApi(token=HF_TOKEN)
 
 
-# ======================================================
-# FREE PREVIEW
-# ======================================================
-def load_free_images():
-    tmp = tempfile.mkdtemp()
-    files = api.list_repo_files(FREE_REPO, token=HF_TOKEN)
+# -------------------------------------------------
+# Helpers
+# -------------------------------------------------
 
-    imgs = [
-        f for f in files
-        if f.startswith(FREE_IMAGE_DIR)
-        and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
-    ]
+def list_date_folders(repo_id):
+    files = api.list_repo_files(repo_id=repo_id)
+    return sorted(
+        {f.split("/")[0] for f in files if "/" in f},
+        reverse=True
+    )
 
-    random.shuffle(imgs)
-    imgs = imgs[:5]
 
-    out = []
-    for f in imgs:
-        try:
-            out.append(hf_hub_download(FREE_REPO, f, token=HF_TOKEN, local_dir=tmp))
-        except Exception:
-            pass
+def list_image_folders(repo_id, date):
+    files = api.list_repo_files(repo_id=repo_id)
+    return sorted(
+        {
+            "/".join(f.split("/")[:2])
+            for f in files
+            if f.startswith(date + "/") and f.count("/") >= 1
+        }
+    )
 
-    return out
 
+def find_image_file(repo_id, folder):
+    files = api.list_repo_files(repo_id=repo_id)
+    for f in files:
+        if f.startswith(folder) and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
+            return f
+    return None
 
-FREE_IMAGES = load_free_images()
 
-# ======================================================
-# HELPERS
-# ======================================================
-def clear_cache():
-    images_cache.clear()
-    prompt_cache.clear()
-    print("[CACHE] cleared")
+def find_summary_file(repo_id, folder):
+    files = api.list_repo_files(repo_id=repo_id)
+    for f in files:
+        if f.startswith(folder) and f.endswith("summary.txt"):
+            return f
+    return None
 
 
-def list_date_folders(repo, last_n=20):
-    files = api.list_repo_files(repo, token=HF_TOKEN)
-    dates = sorted(
-        {f.split("/")[0] for f in files if "/" in f},
-        reverse=True
+def download_file(repo_id, path):
+    return hf_hub_download(
+        repo_id=repo_id,
+        filename=path,
+        token=HF_TOKEN,
+        cache_dir=tempfile.mkdtemp()
     )
-    return dates[:last_n]
 
 
-def download_images(repo, date):
-    key = (repo, date)
-    if key in images_cache:
-        return images_cache[key]
+# -------------------------------------------------
+# Cache (refreshed per repo + date)
+# -------------------------------------------------
 
-    files = api.list_repo_files(repo, token=HF_TOKEN)
-    results = []
+@lru_cache(maxsize=32)
+def load_date(repo_id, date):
+    items = []
 
-    for f in files:
-        if not f.startswith(date + "/"):
-            continue
-        if not f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
+    for folder in list_image_folders(repo_id, date):
+        img_path = find_image_file(repo_id, folder)
+        if not img_path:
             continue
 
-        folder = os.path.dirname(f)
-        dst = os.path.join(tmp_dir, repo.replace("/", "_"), f)
-        os.makedirs(os.path.dirname(dst), exist_ok=True)
-
-        try:
-            img_path = hf_hub_download(
-                repo,
-                f,
-                token=HF_TOKEN,
-                local_dir=os.path.dirname(dst),
-            )
-            prompt = load_prompt_from_folder(repo, folder)
-            results.append((img_path, prompt))
-            print("[DOWNLOAD]", f)
-        except Exception as e:
-            print("[DOWNLOAD ERROR]", f, e)
-
-    images_cache[key] = results
-    print(f"[INFO] {repo} {date}: {len(results)} images")
-    return results
-
-
-# ======================================================
-# PAGINATION
-# ======================================================
-def render_page(repo, date, page):
-    items = download_images(repo, date)
-    total_pages = max(1, math.ceil(len(items) / PAGE_SIZE))
-
-    page = max(0, min(int(page), total_pages - 1))
-    start = page * PAGE_SIZE
-    end = start + PAGE_SIZE
-
-    return (
-        items[start:end],
-        page,
-        f"Page {page+1} / {total_pages}",
-    )
+        local_img = download_file(repo_id, img_path)
 
+        summary_path = find_summary_file(repo_id, folder)
+        summary_text = ""
+        if summary_path:
+            local_summary = download_file(repo_id, summary_path)
+            with open(local_summary, "r", encoding="utf-8") as f:
+                summary_text = f.read().strip()
 
-def nav(delta, repo, date, page):
-    return render_page(repo, date, page + delta)
+        items.append((local_img, summary_text))
 
+    return items
 
-def change_date(repo, date):
-    clear_cache()
-    return render_page(repo, date, 0)
 
+def render_page(repo_id, date, page):
+    all_items = load_date(repo_id, date)
+    total_pages = max(1, (len(all_items) + IMAGES_PER_PAGE - 1) // IMAGES_PER_PAGE)
 
-def change_repo(repo):
-    clear_cache()
-    dates = list_date_folders(repo)
-    first = dates[0] if dates else ""
-    imgs, page, info = render_page(repo, first, 0)
-    return dates, imgs, page, first, info
+    page = max(0, min(page, total_pages - 1))
+    start = page * IMAGES_PER_PAGE
+    end = start + IMAGES_PER_PAGE
 
+    return all_items[start:end], page, total_pages
 
-# ======================================================
-# AUTH
-# ======================================================
-def check_password(pwd):
-    if pwd != APP_PASSWORD:
-        return (
-            gr.update(visible=True),
-            gr.update(visible=False),
-            [],
-            0,
-            "",
-            "❌ Wrong password",
-        )
 
-    dates = list_date_folders(DEFAULT_REPO)
-    latest = dates[0]
+def nav(delta, repo_id, date, page):
+    return (*render_page(repo_id, date, page + delta),)
 
-    imgs, page, info = render_page(DEFAULT_REPO, latest, 0)
 
-    return (
-        gr.update(visible=False),
-        gr.update(visible=True),
-        imgs,
-        page,
-        latest,
-        info,
-    )
+def refresh(repo_id, date):
+    load_date.cache_clear()
+    return render_page(repo_id, date, 0)
 
 
-# ======================================================
-# CSS (UNCHANGED)
-# ======================================================
-css = """
-footer, .gradio-footer { display:none!important; }
-#ohamlab-footer {
-  position:fixed;
-  bottom:0;
-  width:100%;
-  background:#f8f9fb;
-  font-size:12px;
-  padding:8px;
-  border-top:1px solid #e5e7eb;
-  text-align:center;
+# -------------------------------------------------
+# UI
+# -------------------------------------------------
+
+with gr.Blocks(css="""
+footer {visibility: hidden;}
+#copyright {
+  text-align: center;
+  font-size: 12px;
+  color: #888;
+  margin-top: 12px;
 }
-"""
+""") as demo:
 
-# ======================================================
-# UI
-# ======================================================
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    gr.Markdown("# 🔐 OhamLab Image Showcase (PRO)")
-
-    with gr.Column(visible=True) as login_box:
-        pwd = gr.Textbox(type="password", label="Password")
-        login_btn = gr.Button("Unlock")
-
-        gr.Markdown("### 🎁 Free Preview")
-        with gr.Row():
-            for img in FREE_IMAGES:
-                gr.Image(img, show_label=False)
-
-    with gr.Column(visible=False) as app:
-        repo_state = gr.State(DEFAULT_REPO)
-        date_state = gr.State("")
-        page_state = gr.State(0)
-
-        repo_input = gr.Textbox(value=DEFAULT_REPO, label="Base HuggingFace Repo")
-        date_dropdown = gr.Dropdown(label="Select Date")
-        manual_date = gr.Textbox(label="Or enter date (YYYY-MM-DD)")
-
-        gallery = gr.Gallery(columns=3, height="auto", show_label=False)
-        page_info = gr.Markdown()
-
-        with gr.Row():
-            prev_btn = gr.Button("⬅ Previous")
-            next_btn = gr.Button("Next ➡")
-
-    login_btn.click(
-        check_password,
-        inputs=pwd,
-        outputs=[login_box, app, gallery, page_state, date_state, page_info],
+    gr.Markdown("## 📸 Image Browser (Private HF Repo)")
+
+    with gr.Row():
+        repo_input = gr.Textbox(
+            value=DEFAULT_REPO,
+            label="Repository ID"
+        )
+        date_dd = gr.Dropdown(
+            choices=list_date_folders(DEFAULT_REPO),
+            label="Date"
+        )
+
+    gallery = gr.Gallery(
+        label="Images",
+        columns=3,
+        height="auto",
+        show_label=True
     )
 
-    repo_input.change(
-        change_repo,
-        inputs=repo_input,
-        outputs=[date_dropdown, gallery, page_state, date_state, page_info],
-    ).then(lambda r: r, inputs=repo_input, outputs=repo_state)
+    page_state = gr.State(0)
+    page_info = gr.Markdown()
 
-    date_dropdown.change(
-        lambda r, d: change_date(r, d),
-        inputs=[repo_state, date_dropdown],
-        outputs=[gallery, page_state, page_info],
-    ).then(lambda d: d, inputs=date_dropdown, outputs=date_state)
+    with gr.Row():
+        prev_btn = gr.Button("⬅ Prev")
+        next_btn = gr.Button("Next ➡")
+        refresh_btn = gr.Button("🔄 Refresh")
 
-    manual_date.submit(
-        lambda r, d: change_date(r, d),
-        inputs=[repo_state, manual_date],
-        outputs=[gallery, page_state, page_info],
-    ).then(lambda d: d, inputs=manual_date, outputs=date_state)
+    gr.HTML("<div id='copyright'>© OhamLab</div>")
 
+    # Initial load
+    demo.load(
+        lambda r, d: render_page(r, d, 0),
+        inputs=[repo_input, date_dd],
+        outputs=[gallery, page_state, page_info]
+    )
+
+    # Navigation
     prev_btn.click(
         lambda r, d, p: nav(-1, r, d, p),
-        inputs=[repo_state, date_state, page_state],
-        outputs=[gallery, page_state, page_info],
+        inputs=[repo_input, date_dd, page_state],
+        outputs=[gallery, page_state, page_info]
     )
 
     next_btn.click(
         lambda r, d, p: nav(1, r, d, p),
-        inputs=[repo_state, date_state, page_state],
-        outputs=[gallery, page_state, page_info],
+        inputs=[repo_input, date_dd, page_state],
+        outputs=[gallery, page_state, page_info]
     )
 
-    gr.HTML("<div id='ohamlab-footer'>© OhamLab Copyright</div>")
+    refresh_btn.click(
+        refresh,
+        inputs=[repo_input, date_dd],
+        outputs=[gallery, page_state, page_info]
+    )
+
+    # Change date or repo → refresh cache
+    date_dd.change(
+        refresh,
+        inputs=[repo_input, date_dd],
+        outputs=[gallery, page_state, page_info]
+    )
+
+    repo_input.change(
+        lambda r: gr.Dropdown(choices=list_date_folders(r)),
+        inputs=repo_input,
+        outputs=date_dd
+    )
 
 demo.launch()