Update app.py

app.py

@@ -552,18 +552,19 @@ async def process_dataset_task(start_index: int):
                 # Final progress report for the current file
                 if len(all_captions) == len(image_paths):
                     print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: Successfully completed all {len(all_captions)} captions.")
-                    current_file_success = True
                 else:
-                    print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: Completed with partial result: {len(all_captions)}/{len(image_paths)} captions. Marking as partial failure.")
-                    current_file_success = False
+                    print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: Completed with partial result: {len(all_captions)}/{len(image_paths)} captions.")
+                
+                # Consider the file successful if we have any captions
+                current_file_success = len(all_captions) > 0
                 
                 # 4. Upload Results
                 if all_captions:
                     print(f"[{FLOW_ID}] Uploading {len(all_captions)} captions for {zip_full_name}...")
                     if await upload_captions_to_hf(zip_full_name, all_captions):
                         print(f"[{FLOW_ID}] Successfully uploaded captions for {zip_full_name}.")
-                        # Partial success in captioning is still a success for the upload step
-                        pass
+                        # Keep current_file_success as True since we have captions and successfully uploaded them
+                        current_file_success = True
                     else:
                         print(f"[{FLOW_ID}] Failed to upload captions for {zip_full_name}.")
                         current_file_success = False
@@ -615,12 +616,22 @@ app = FastAPI(
 async def startup_event():
     print(f"Flow Server {FLOW_ID} started on port {FLOW_PORT}.")
     
-    # Automatically start the processing task
+    # Get both local progress and HF state
     progress = load_progress()
-    # Start from the last processed index + 1, or the hardcoded AUTO_START_INDEX if the progress file is new/empty
-    start_index = progress.get('last_processed_index', 0) + 1
-    if start_index < AUTO_START_INDEX:
-        start_index = AUTO_START_INDEX
+    current_state = await download_hf_state()
+    
+    # Get the next_download_index from HF state if available
+    hf_next_index = current_state.get("next_download_index", 0)
+    
+    # If HF state has a higher index, use that instead of local progress
+    if hf_next_index > 0:
+        start_index = hf_next_index
+        print(f"[{FLOW_ID}] Using next_download_index from HF state: {start_index}")
+    else:
+        # Fall back to local progress if HF state doesn't have a meaningful index
+        start_index = progress.get('last_processed_index', 0) + 1
+        if start_index < AUTO_START_INDEX:
+            start_index = AUTO_START_INDEX
 
     # Use a dummy BackgroundTasks object for the startup task
     # Note: FastAPI's startup events can't directly use BackgroundTasks, but we can use asyncio.create_task