Update app.py

app.py

@@ -34,28 +34,29 @@ def quota_guard(fn):
         return fn(*args)
     return wrapper
 
-
 # ======================================================
 # CONFIG
 # ======================================================
-REPO_ID = "rahul7star/WAN22-Mar-2026"  # PRO repo
+REPO_ID = "rahul7star/WAN22-Mar-2026"
 FREE_REPO = "rahul7star/ohamlab"
 FREE_IMAGE_DIR = "showcase/image"
+HF_TOKEN = os.getenv("HF_TOKEN")  # Must have read access to private repo
+
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN env var not set")
 
-HF_TOKEN = os.getenv("HF_TOKEN")  # HuggingFace token for private repo
 APP_PASSWORD = os.getenv("APP_PASSWORD")
 if not APP_PASSWORD:
     raise RuntimeError("APP_PASSWORD env var not set")
 
 api = HfApi()
 
-
 # ======================================================
 # FREE IMAGE PREVIEW
 # ======================================================
 def load_free_images():
     tmp = tempfile.mkdtemp()
-    files = api.list_repo_files(FREE_REPO)
+    files = api.list_repo_files(FREE_REPO, token=HF_TOKEN)
 
     images = [
         f for f in files
@@ -72,74 +73,76 @@ def load_free_images():
     local = []
     for img in picked:
         try:
-            local.append(hf_hub_download(FREE_REPO, img, local_dir=tmp))
-        except Exception:
-            pass
+            local_path = hf_hub_download(FREE_REPO, img, local_dir=tmp, token=HF_TOKEN)
+            local.append(local_path)
+        except Exception as e:
+            print(f"[FREE IMG ERROR] {img}: {e}")
 
     return local
 
 FREE_IMAGES = load_free_images()
 
-
 # ======================================================
 # PRO HELPERS
 # ======================================================
-def list_date_folders():
-    """Return last 20 dates from repo"""
+tmp_dir = tempfile.mkdtemp()
+images_cache = {}  # date_folder -> list of local image paths
+
+def list_date_folders(last_n_days=20):
     files = api.list_repo_files(REPO_ID, token=HF_TOKEN)
     dates = sorted({f.split("/")[0] for f in files if "/" in f}, reverse=True)
-    return dates[:20]
+    return dates[:last_n_days]
 
+def download_images_for_date(date_folder):
+    if date_folder in images_cache:
+        return images_cache[date_folder]
 
-def load_images_for_date(date_folder):
-    """Download images locally and return local paths"""
-    tmp = tempfile.mkdtemp()
     files = api.list_repo_files(REPO_ID, token=HF_TOKEN)
-    images = []
+    local_paths = []
 
     for f in files:
         if f.startswith(date_folder + "/") and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
-            # always save as input_image.png to avoid weird names
-            local_path = os.path.join(tmp, *f.split("/"))
-            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+            sub_path = os.path.join(tmp_dir, f)
+            os.makedirs(os.path.dirname(sub_path), exist_ok=True)
             try:
-                hf_hub_download(REPO_ID, f, local_dir=os.path.dirname(local_path), token=HF_TOKEN, local_file_name=os.path.basename(local_path))
-                images.append(local_path)
-                print(f"[DOWNLOAD] {f} -> {local_path}")
+                local_path = hf_hub_download(
+                    REPO_ID,
+                    f,
+                    local_dir=os.path.dirname(sub_path),
+                    token=HF_TOKEN
+                )
+                # Rename to a consistent name
+                final_path = os.path.join(os.path.dirname(local_path), "input_image.png")
+                os.rename(local_path, final_path)
+                local_paths.append(final_path)
+                print(f"[DOWNLOAD] {f} -> {final_path}")
             except Exception as e:
-                print(f"[ERROR] Download failed {f}: {e}")
-
-    print(f"[INFO] Total images for {date_folder}: {len(images)}")
-    return images
+                print(f"[DOWNLOAD ERROR] {f}: {e}")
 
+    images_cache[date_folder] = local_paths
+    print(f"[INFO] Total images for {date_folder}: {len(local_paths)}")
+    return local_paths
 
 # ======================================================
 # PAGINATION
 # ======================================================
-def render_grid(date_folder, page, page_size):
-    data = load_images_for_date(date_folder)
-    total_pages = max(1, math.ceil(len(data) / page_size))
-
-    page = max(0, min(int(page), total_pages - 1))
+def render_page(date_folder, page, page_size):
+    images = download_images_for_date(date_folder)
+    total_pages = max(1, math.ceil(len(images) / page_size))
+    page = max(0, min(page, total_pages - 1))
     start, end = page * page_size, (page + 1) * page_size
-    page_images = data[start:end]
-
-    print(f"[RENDER] Page {page+1}/{total_pages}, images: {page_images}")
-
-    return page_images, page
-
-
-def next_page(date, page, size):
-    return render_grid(date, page + 1, size)
+    subset = images[start:end]
+    print(f"[RENDER] Page {page+1}/{total_pages}, images: {subset}")
+    return subset, page
 
+def next_page(date_folder, page, page_size):
+    return render_page(date_folder, page + 1, page_size)
 
-def prev_page(date, page, size):
-    return render_grid(date, page - 1, size)
-
-
-def reset_page(date, size):
-    return render_grid(date, 0, size)
+def prev_page(date_folder, page, page_size):
+    return render_page(date_folder, page - 1, page_size)
 
+def reset_page(date_folder, page_size):
+    return render_page(date_folder, 0, page_size)
 
 # ======================================================
 # AUTH
@@ -151,15 +154,14 @@ def check_password(pwd):
             gr.update(visible=True),
             gr.update(visible=False),
             gr.update(visible=True),
-            [],
-            0,
+            [], 0,
             "❌ Wrong password",
             ""
         )
 
     dates = list_date_folders()
     latest = dates[0] if dates else ""
-    images, page = render_grid(latest, 0, 6)
+    images, page = render_page(latest, 0, 6)
 
     return (
         gr.update(visible=False),
@@ -171,13 +173,25 @@ def check_password(pwd):
         latest
     )
 
-
 # ======================================================
 # CSS
 # ======================================================
 css = """
+.grid {
+  display:grid;
+  grid-template-columns:repeat(auto-fill,minmax(220px,1fr));
+  gap:16px;
+}
+.card {
+  background:#0b1220;
+  padding:10px;
+  border-radius:14px;
+}
+.card img {
+  width:100%;
+  border-radius:10px;
+}
 footer, .gradio-footer { display:none!important; }
-
 #ohamlab-footer {
   position:fixed;
   bottom:0;
@@ -190,7 +204,6 @@ footer, .gradio-footer { display:none!important; }
 }
 """
 
-
 # ======================================================
 # UI
 # ======================================================
@@ -214,13 +227,23 @@ with gr.Blocks() as demo:
         current_date = gr.State("")
         page_state = gr.State(0)
         page_size = gr.State(6)
+        images_state = gr.State([])
 
+        # Gallery to render images
         gallery = gr.Gallery(label="Images", elem_id="gallery", columns=3, height="auto")
 
+        page_info = gr.Markdown()
+
+        # Navigation buttons
         with gr.Row():
             prev_btn = gr.Button("⬅ Previous")
             next_btn = gr.Button("Next ➡")
 
+        def update_gallery(images, page, page_size):
+            subset, page = render_page(current_date.value, page, page_size)
+            page_info.update(f"Page {page+1} / {math.ceil(len(images_state.value)/page_size)}")
+            return subset, page
+
         prev_btn.click(prev_page, [current_date, page_state, page_size], [gallery, page_state])
         next_btn.click(next_page, [current_date, page_state, page_size], [gallery, page_state])
 
@@ -232,4 +255,4 @@ with gr.Blocks() as demo:
 
     gr.HTML("<div id='ohamlab-footer'>© OhamLab Copyright</div>")
 
-demo.launch(share=False)
+demo.launch()