Update app.py

app.py

@@ -7,8 +7,8 @@ import asyncio
 import threading
 import time
 import hashlib
-from typing import Dict, List, Set
-from fastapi import FastAPI, BackgroundTasks, HTTPException
+from typing import Dict, List, Set, Optional
+from fastapi import FastAPI, BackgroundTasks, HTTPException, Form
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import HfApi, list_repo_files
@@ -101,7 +101,7 @@ def load_download_state() -> int:
     if os.path.exists(DOWNLOAD_STATE_FILE):
         try:
             with open(DOWNLOAD_STATE_FILE, "r") as f:
-                return json.load(f).get("next_download_index", 3)
+                return json.load(f).get("next_download_index", 0)
         except json.JSONDecodeError:
             log_message(f"⚠️ Warning: Could not decode {DOWNLOAD_STATE_FILE}. Starting download from index 0.")
             return 0
@@ -129,7 +129,7 @@ def save_processed_files_state(processed_set: set):
     with open(PROCESS_STATE_FILE, "w") as f:
         json.dump({"processed_rars": list(processed_set)}, f)
 
-def download_rar_files(start_index: 4, chunk_size: int) -> tuple:
+def download_rar_files(start_index: int, chunk_size: int) -> tuple:
     """Downloads a batch of RAR files from the source dataset"""
     try:
         all_files = list_repo_files(repo_id=SOURCE_REPO_ID, repo_type="dataset", token=HF_TOKEN)
@@ -143,7 +143,7 @@ def download_rar_files(start_index: 4, chunk_size: int) -> tuple:
             log_message("✅ No more RAR files to download.")
             return [], start_index
 
-        log_message(f"📥 Downloading RAR files {start_index + 4} to {end_index} from {SOURCE_REPO_ID}")
+        log_message(f"📥 Downloading RAR files {start_index + 1} to {end_index} from {SOURCE_REPO_ID}")
 
         downloaded_paths = []
         for file_path_in_repo in files_to_download_metadata:
@@ -190,7 +190,7 @@ def extract_and_upload_rar(rar_path: str, processed_rars_set: set, uploaded_fold
     
     # Check if this folder content has already been uploaded
     if folder_hash in uploaded_folders_set:
-        log_message(f"🔒 Folder '{folder_name}' already uploaded (hash: {folder_hash[:8]}...), skipping.")
+        log_message(f"🔒 Folder \'{folder_name}\' already uploaded (hash: {folder_hash[:8]}...), skipping.")
         processed_rars_set.add(filename)
         save_processed_files_state(processed_rars_set)
         return True
@@ -251,7 +251,7 @@ def extract_and_upload_rar(rar_path: str, processed_rars_set: set, uploaded_fold
         save_uploaded_folders(uploaded_folders_set)
         processing_status["uploaded_folders"] = len(uploaded_folders_set)
         
-        log_message(f"🔒 Folder '{folder_name}' locked in BG2 repo (hash: {folder_hash[:8]}...)")
+        log_message(f"🔒 Folder \'{folder_name}\' locked in BG2 repo (hash: {folder_hash[:8]}...)")
         
         return True
 
@@ -275,8 +275,8 @@ def extract_and_upload_rar(rar_path: str, processed_rars_set: set, uploaded_fold
             except Exception as e:
                 log_message(f"⚠️ Could not clean up extraction folder {current_extract_folder}: {e}")
 
-def continuous_processing():
-    """Main processing loop that runs continuously"""
+def continuous_processing(start_download_index: Optional[int] = None):
+    """Main processing loop that runs continuously, with an optional starting download index"""
     processing_status["is_running"] = True
     log_message("🚀 Starting continuous RAR processing...")
     
@@ -285,6 +285,12 @@ def continuous_processing():
         uploaded_folders = load_uploaded_folders()
         processing_status["uploaded_folders"] = len(uploaded_folders)
         
+        if start_download_index is not None:
+            log_message(f"Starting download from index: {start_download_index}")
+            save_download_state(start_download_index) # Set download state to start from this index
+        else:
+            log_message("Starting download from saved state or beginning.")
+
         while processing_status["is_running"]:
             # 1. Download a batch of RAR files
             download_start_index = load_download_state()
@@ -298,7 +304,7 @@ def continuous_processing():
                 unprocessed_rars = [rar for rar in all_local_rars if os.path.basename(rar) not in processed_rars]
                 
                 if not unprocessed_rars:
-                    log_message("✅ All RAR files have been processed. Stopping...")
+                    log_message("✅ No more RAR files to download. Stopping...")
                     break
                 else:
                     log_message(f"📋 Found {len(unprocessed_rars)} unprocessed local RAR files")
@@ -365,6 +371,10 @@ async def root():
             .stop-button:hover { background: #d32f2f; }
             .stats { display: flex; gap: 20px; margin: 20px 0; }
             .stat-item { background: #f0f0f0; padding: 10px; border-radius: 5px; text-align: center; flex: 1; }
+            .start-form { margin-top: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 5px; background: #f9f9f9; }
+            .start-form input[type="number"] { width: calc(100% - 120px); padding: 8px; margin-right: 10px; border: 1px solid #ccc; border-radius: 4px; }
+            .start-form button { padding: 8px 15px; background: #4CAF50; color: white; border: none; border-radius: 4px; cursor: pointer; }
+            .start-form button:hover { background: #45a049; }
         </style>
     </head>
     <body>
@@ -397,8 +407,14 @@ async def root():
                 </div>
             </div>
             
+            <div class="start-form">
+                <h3>Start Processing from Specific Download Index</h3>
+                <input type="number" id="start-index-input" placeholder="Enter start index (e.g., 0)" value="0">
+                <button onclick="startProcessingWithIndex()">Start from Index</button>
+            </div>
+
             <div>
-                <button class="button" onclick="startProcessing()" id="start-btn">Start Processing</button>
+                <button class="button" onclick="startProcessing()" id="start-btn">Start Processing (from last saved index)</button>
                 <button class="button stop-button" onclick="stopProcessing()" id="stop-btn" disabled>Stop Processing</button>
                 <button class="button" onclick="refreshStatus()">Refresh Status</button>
             </div>
@@ -410,47 +426,67 @@ async def root():
         <script>
             async function startProcessing() {
                 try {
-                    const response = await fetch('/start', { method: 'POST' });
+                    const response = await fetch(\"/start\", { method: \"POST\" });
+                    const result = await response.json();
+                    alert(result.message);
+                    refreshStatus();
+                } catch (error) {
+                    alert(\"Error starting processing: \" + error.message);
+                }
+            }
+
+            async function startProcessingWithIndex() {
+                const index = document.getElementById(\"start-index-input\").value;
+                if (index === "" || isNaN(index)) {
+                    alert(\"Please enter a valid number for the start index.\");
+                    return;
+                }
+                try {
+                    const response = await fetch(\"/start_from_index\", {
+                        method: \"POST\",
+                        headers: { \"Content-Type\": \"application/x-www-form-urlencoded\" },
+                        body: `start_index=${parseInt(index)}`
+                    });
                     const result = await response.json();
                     alert(result.message);
                     refreshStatus();
                 } catch (error) {
-                    alert('Error starting processing: ' + error.message);
+                    alert(\"Error starting processing from index: \" + error.message);
                 }
             }
             
             async function stopProcessing() {
                 try {
-                    const response = await fetch('/stop', { method: 'POST' });
+                    const response = await fetch(\"/stop\", { method: \"POST\" });
                     const result = await response.json();
                     alert(result.message);
                     refreshStatus();
                 } catch (error) {
-                    alert('Error stopping processing: ' + error.message);
+                    alert(\"Error stopping processing: \" + error.message);
                 }
             }
             
             async function refreshStatus() {
                 try {
-                    const response = await fetch('/status');
+                    const response = await fetch(\"/status\");
                     const status = await response.json();
                     
-                    document.getElementById('status').textContent = status.is_running ? 'Running' : 'Stopped';
-                    document.getElementById('current-file').textContent = status.current_file || 'None';
-                    document.getElementById('last-update').textContent = status.last_update || 'Never';
-                    document.getElementById('total-files').textContent = status.total_files;
-                    document.getElementById('processed-files').textContent = status.processed_files;
-                    document.getElementById('uploaded-folders').textContent = status.uploaded_folders;
-                    document.getElementById('failed-files').textContent = status.failed_files;
+                    document.getElementById(\"status\").textContent = status.is_running ? \"Running\" : \"Stopped\";
+                    document.getElementById(\"current-file\").textContent = status.current_file || \"None\";
+                    document.getElementById(\"last-update\").textContent = status.last_update || \"Never\";
+                    document.getElementById(\"total-files\").textContent = status.total_files;
+                    document.getElementById(\"processed-files\").textContent = status.processed_files;
+                    document.getElementById(\"uploaded-folders\").textContent = status.uploaded_folders;
+                    document.getElementById(\"failed-files\").textContent = status.failed_files;
                     
-                    document.getElementById('start-btn').disabled = status.is_running;
-                    document.getElementById('stop-btn').disabled = !status.is_running;
+                    document.getElementById(\"start-btn\").disabled = status.is_running;
+                    document.getElementById(\"stop-btn\").disabled = !status.is_running;
                     
-                    const logsDiv = document.getElementById('logs');
-                    logsDiv.innerHTML = status.logs.join('<br>');
+                    const logsDiv = document.getElementById(\"logs\");
+                    logsDiv.innerHTML = status.logs.join(\"<br>\");
                     logsDiv.scrollTop = logsDiv.scrollHeight;
                 } catch (error) {
-                    console.error('Error refreshing status:', error);
+                    console.error(\"Error refreshing status:\", error);
                 }
             }
             
@@ -479,6 +515,18 @@ async def start_processing(background_tasks: BackgroundTasks):
     background_tasks.add_task(continuous_processing)
     return {"message": "Processing started"}
 
+@app.post("/start_from_index")
+async def start_processing_from_index(background_tasks: BackgroundTasks, start_index: int = Form(...)):
+    """Start the processing from a specific download index in background"""
+    if processing_status["is_running"]:
+        return {"message": "Processing is already running"}
+    
+    if start_index < 0:
+        return {"message": "Start index cannot be negative."}
+
+    background_tasks.add_task(continuous_processing, start_download_index=start_index)
+    return {"message": f"Processing started from download index: {start_index}"}
+
 @app.post("/stop")
 async def stop_processing():
     """Stop the processing"""
@@ -502,3 +550,4 @@ async def get_uploaded_folders():
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)
 
+