Update app.py

app.py

@@ -58,6 +58,7 @@ FFMPEG_CANDIDATES = [
 ]
 MODEL_CACHE = {}
 EXTRACT_MAP = {}  # friendly_name -> path
+LAST_EXTRACT_DIR = None  # path to last extraction folder (for download)
 DEFAULT_ZIP_PASS = "dietcoke1"
 
 CPU_COUNT = max(1, multiprocessing.cpu_count())
@@ -464,8 +465,14 @@ def _worker_transcribe(args):
 
 # ---------- ZIP extraction & mapping ----------
 def extract_zip_and_map(zip_path, zip_password=None):
-    global EXTRACT_MAP
+    """
+    Extract ZIP into a per-run temp dir, populate EXTRACT_MAP (friendly name -> file path),
+    and set LAST_EXTRACT_DIR to the extraction folder for download.
+    Returns (friendly_list, logs_str)
+    """
+    global EXTRACT_MAP, LAST_EXTRACT_DIR
     EXTRACT_MAP = {}
+    LAST_EXTRACT_DIR = None
     run_id = uuid4().hex
     temp_extract_dir = os.path.join(tempfile.gettempdir(), f"extracted_audio_{run_id}")
     logs = []
@@ -509,8 +516,15 @@ def extract_zip_and_map(zip_path, zip_password=None):
                 logs.append(f"Extracted: {info.filename} -> {key}")
         if not EXTRACT_MAP:
             logs.append("No supported audio files found in ZIP.")
+            # cleanup temp dir if empty
+            try:
+                if os.path.exists(temp_extract_dir) and not os.listdir(temp_extract_dir):
+                    shutil.rmtree(temp_extract_dir)
+            except Exception:
+                pass
             return [], "\n".join(logs)
         friendly = sorted(EXTRACT_MAP.keys())
+        LAST_EXTRACT_DIR = temp_extract_dir
         return friendly, "\n".join(logs)
     except Exception as e:
         traceback.print_exc()
@@ -521,6 +535,27 @@ def extract_zip_and_map(zip_path, zip_password=None):
             pass
         return [], f"Extraction failed: {e}"
 
+def download_extracted_folder():
+    """
+    Zip LAST_EXTRACT_DIR and return zip path for download (or None + message if missing).
+    """
+    global LAST_EXTRACT_DIR
+    if not LAST_EXTRACT_DIR or not os.path.exists(LAST_EXTRACT_DIR):
+        return None, "No extracted folder available for download."
+    try:
+        zip_tmp = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+        zip_tmp.close()
+        with zipfile.ZipFile(zip_tmp.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+            # Walk and add files preserving relative path
+            for root, dirs, files in os.walk(LAST_EXTRACT_DIR):
+                for f in files:
+                    fullp = os.path.join(root, f)
+                    rel = os.path.relpath(fullp, LAST_EXTRACT_DIR)
+                    zf.write(fullp, arcname=rel)
+        return zip_tmp.name, "OK"
+    except Exception as e:
+        return None, f"Failed to create ZIP: {e}"
+
 # ---------- Batch transcription generator (streaming) ----------
 def batch_transcribe_parallel_generator(
     friendly_selected,
@@ -883,9 +918,12 @@ with gr.Blocks(title="Whisper Transcriber — Parallel + Memory", css=CSS) as de
                     batch_files = gr.File(label="Upload audio files", file_count="multiple", type="filepath")
                     batch_zip = gr.File(label="Or upload ZIP (optional)", file_count="single", type="filepath")
                     batch_zip_pass = gr.Textbox(label="ZIP password (if any)", value=DEFAULT_ZIP_PASS)
-                    batch_preview_btn = gr.Button("Preview ZIP members")
+                    batch_preview_btn = gr.Button("Extract & List ZIP files")
                     batch_preview_out = gr.Textbox(label="ZIP members (preview)", lines=6, interactive=False)
                     batch_select = gr.CheckboxGroup(choices=[], label="Select extracted files to include", interactive=True)
+                    # NEW: download extracted folder button + output file
+                    batch_download_extracted_btn = gr.Button("Download extracted folder")
+                    batch_extracted_zip = gr.File(label="Downloaded extracted ZIP")
                     batch_model = gr.Dropdown(choices=available_choices, value=default_choice, label="Model")
                     batch_device = gr.Dropdown(choices=["auto", "cpu", "cuda"], value="auto", label="Device")
                     batch_merge = gr.Checkbox(label="Merge transcripts into DOCX", value=True)
@@ -902,16 +940,30 @@ with gr.Blocks(title="Whisper Transcriber — Parallel + Memory", css=CSS) as de
                     batch_zip_download = gr.File(label="Download per-file transcripts ZIP")
                     batch_doc_download = gr.File(label="Download merged DOCX (if created)")
 
-            def _preview_zip(zip_file, password):
+            def _preview_zip_and_populate(zip_file, password):
+                """
+                Extract the zip, populate EXTRACT_MAP and return updated CheckboxGroup choices + logs string.
+                """
                 if not zip_file:
-                    return "No ZIP provided."
+                    return gr.update(choices=[]), "No ZIP provided."
                 path = zip_file.name if hasattr(zip_file, "name") else str(zip_file)
-                members, logs = extract_zip_and_map(path, password)
-                if members:
-                    return "\n".join(members)
-                return logs
-
-            batch_preview_btn.click(fn=_preview_zip, inputs=[batch_zip, batch_zip_pass], outputs=[batch_preview_out])
+                friendly, logs = extract_zip_and_map(path, password)
+                if friendly:
+                    # return an update for the CheckboxGroup (populates choices list) and a logs string
+                    return gr.update(choices=friendly), "\n".join(logs.splitlines())
+                return gr.update(choices=[]), logs
+
+            # wire preview/extract button to update the batch_select choices and preview textbox
+            batch_preview_btn.click(fn=_preview_zip_and_populate, inputs=[batch_zip, batch_zip_pass], outputs=[batch_select, batch_preview_out])
+
+            def _download_extracted_wrapper():
+                zip_path, msg = download_extracted_folder()
+                if zip_path:
+                    return zip_path
+                # gr.File expects path or None; if failed, return None so nothing is downloadable
+                return None
+
+            batch_download_extracted_btn.click(fn=_download_extracted_wrapper, inputs=[], outputs=[batch_extracted_zip])
 
             # wrapper generator — Gradio expects the function itself to be a generator that yields streaming tuples
             def _start_batch(friendly_selected, uploaded_files, zip_file, zip_pass, model_name, device_name, merge_flag, mem_flag, srt_flag, use_two_pass, fast_model, refine_thresh):