Update app.py

app.py

@@ -348,9 +348,9 @@ def process_file_background(task_id: str, file_data: bytes, filename: str, extra
         output_dir = Path(app.config['OUTPUT_FOLDER']) / stem
         output_dir.mkdir(parents=True, exist_ok=True)
         
-        # Copy PDF to output directory
+        # Copy PDF to output directory, using replace to overwrite if it already exists
         pdf_path = output_dir / filename
-        upload_path.rename(pdf_path)
+        upload_path.replace(pdf_path)
         
         _update_task_progress(task_id, filename, 25, f'Loading model and processing {filename}...')
         
@@ -812,6 +812,48 @@ def delete_pdf_by_path(stem: str):
         return jsonify({'error': str(e)}), 500
 
 
+@app.route('/api/download-zip/<path:stem>', methods=['GET'])
+def download_zip(stem: str):
+    """Download the processed output as a zip archive."""
+    import io
+    import zipfile
+    
+    stem = stem.strip()
+    if not stem:
+        return jsonify({'error': 'Missing stem'}), 400
+        
+    output_root = Path(app.config['OUTPUT_FOLDER']).resolve()
+    target_dir = (output_root / stem).resolve()
+    
+    # Prevent path traversal
+    if output_root not in target_dir.parents and target_dir != output_root:
+        return jsonify({'error': 'Invalid stem path'}), 400
+        
+    if not target_dir.exists() or not target_dir.is_dir():
+        return jsonify({'error': 'PDF not found or not processed completely'}), 404
+        
+    try:
+        memory_file = io.BytesIO()
+        with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
+            for root, _, files in os.walk(target_dir):
+                for file in files:
+                    file_path = Path(root) / file
+                    arcname = file_path.relative_to(target_dir)
+                    zf.write(file_path, arcname)
+                    
+        memory_file.seek(0)
+        
+        return send_file(
+            memory_file,
+            mimetype='application/zip',
+            as_attachment=True,
+            download_name=f"{stem}_extracted.zip"
+        )
+    except Exception as e:
+        logger.error(f"Zip creation failed: {e}")
+        return jsonify({'error': str(e)}), 500
+
+
 if __name__ == '__main__':
     # Run on port 7860 for Hugging Face Spaces, or 5000 for local development
     port = int(os.environ.get('PORT', 7860))