Update app.py

app.py

@@ -5,24 +5,29 @@ import time
 import gradio as gr
 from gradio_client import Client, handle_file
 
+# -------------------------
+# CONFIG
+# -------------------------
+
 SPACE = "pockit-cloud/main"
 client = Client(SPACE)
 
 SPACE_URL = os.getenv("SPACE_URL", "https://pockit-cloud-main.hf.space")
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")  # Secret for private dataset access
 
 TMP_DIR = "tmp"
 os.makedirs(TMP_DIR, exist_ok=True)
 
 # -------------------------
-# CLEAN OLD TEMP FILES
+# CLEAN TEMP FILES
 # -------------------------
 def cleanup_tmp():
     now = time.time()
     for f in os.listdir(TMP_DIR):
         path = os.path.join(TMP_DIR, f)
-        if os.path.isfile(path) and now - os.path.getmtime(path) > 3600:
-            os.remove(path)
+        if os.path.isfile(path):
+            if now - os.path.getmtime(path) > 3600:  # 1 hour
+                os.remove(path)
 
 # -------------------------
 # LOGIN
@@ -59,7 +64,7 @@ def list_files(user_id, password):
         return str(e), []
 
 # -------------------------
-# UPLOAD FILE
+# UPLOAD
 # -------------------------
 def upload_file(user_id, password, file, custom_name):
 
@@ -82,19 +87,21 @@ def upload_file(user_id, password, file, custom_name):
         return str(e)
 
 # -------------------------
-# DOWNLOAD USING WGET MIRROR
+# DOWNLOAD (PRIVATE SAFE VERSION)
+# Returns file directly to gr.File
 # -------------------------
 def download_file(user_id, password, filename):
 
     if not user_id or not password:
-        return "Username and password required"
+        return None
 
     if not filename:
-        return "No file selected"
+        return None
 
     try:
         cleanup_tmp()
 
+        # Get private download link
         link = client.predict(
             user_id=user_id,
             password=password,
@@ -114,18 +121,20 @@ def download_file(user_id, password, filename):
             link
         ]
 
+        # Add HuggingFace token if available
         if HF_TOKEN:
             cmd.insert(1, f"--header=Authorization: Bearer {HF_TOKEN}")
 
         subprocess.run(cmd, check=True)
 
-        return f"{SPACE_URL}/file={temp_path}"
+        return temp_path  # IMPORTANT: return path to gr.File
 
     except Exception as e:
-        return str(e)
+        print(e)
+        return None
 
 # -------------------------
-# DELETE FILE
+# DELETE
 # -------------------------
 def delete_file(user_id, password, filename):
 
@@ -197,10 +206,10 @@ with gr.Blocks(title="Pockit Cloud Client") as app:
             outputs=[status_box, file_list]
         )
 
-        gr.Markdown("### Download File")
+        gr.Markdown("### Download")
 
-        download_btn = gr.Button("Download")
-        download_output = gr.Textbox(label="Download Link")
+        download_btn = gr.Button("Download File")
+        download_output = gr.File(label="Downloaded File")
 
         download_btn.click(
             download_file,
@@ -208,7 +217,7 @@ with gr.Blocks(title="Pockit Cloud Client") as app:
             outputs=download_output
         )
 
-        gr.Markdown("### Upload File")
+        gr.Markdown("### Upload")
 
         upload_input = gr.File(label="Select File")
         custom_name = gr.Textbox(label="Custom Name (Optional)")
@@ -221,7 +230,7 @@ with gr.Blocks(title="Pockit Cloud Client") as app:
             outputs=upload_status
         )
 
-        gr.Markdown("### Delete File")
+        gr.Markdown("### Delete")
 
         delete_btn = gr.Button("Delete Selected File")
         delete_status = gr.Textbox(label="Delete Status")
@@ -243,5 +252,5 @@ with gr.Blocks(title="Pockit Cloud Client") as app:
             outputs=change_status
         )
 
-# IMPORTANT: allow tmp folder
+# IMPORTANT: allow tmp folder for serving
 app.launch(allowed_paths=[TMP_DIR])