Update app.py

app.py

@@ -348,26 +348,37 @@ async def upload_captions_to_hf(zip_full_name: str, captions: List[Dict]) -> boo
         return False
 
 async def process_next_file_task():
-    """Task to process the next file in the list based on the current index."""
+    """Continuous task to process files based on the current index."""
     global state
     
-    if not state.is_running:
-        print(f"[{FLOW_ID}] Processing loop is not running. Exiting task.")
-        return
-        
-    while state.is_running:
-        repo_file_full_path = None
-        current_index = -1
-        
+    print(f"[{FLOW_ID}] Processing task started. Running: {state.is_running}")
+    
+    while True:
+        # Check if we should be running
         async with state_lock:
+            should_run = state.is_running
             current_index = state.current_index
-            if current_index >= state.total_files:
+        
+        if not should_run:
+            # Wait a bit and check again
+            await asyncio.sleep(2)
+            continue
+            
+        # Check if we have files to process
+        if current_index >= state.total_files:
+            async with state_lock:
                 state.status = "Finished processing all files."
                 state.is_running = False
-                print(f"[{FLOW_ID}] Reached end of file list. Stopping processing.")
-                await save_state_to_hf()
-                break
-                
+                state.current_file = None
+                state.current_file_progress = "0/0"
+            print(f"[{FLOW_ID}] Reached end of file list. Stopping processing.")
+            await save_state_to_hf()
+            await asyncio.sleep(2)
+            continue
+        
+        # Process the current file
+        repo_file_full_path = None
+        async with state_lock:
             repo_file_full_path = state.all_zip_files[current_index]
             
             if repo_file_full_path in state.processed_files:
@@ -378,18 +389,19 @@ async def process_next_file_task():
                 await save_state_to_hf()
                 continue
                 
-            # Mark the file as in-progress in the state
+            # Mark the file as in-progress
             state.status = f"Processing file {current_index + 1}/{state.total_files}"
             state.current_file = Path(repo_file_full_path).name
             state.current_file_progress = "0/0"
             await save_state_to_hf()
-            
-        # --- Start Processing ---
+        
+        # --- Process the file ---
         extract_dir = None
         zip_full_name = None
-        global_success = False
         
         try:
+            print(f"[{FLOW_ID}] Processing: {repo_file_full_path}")
+            
             download_result = await download_and_extract_zip(repo_file_full_path)
             
             if download_result is None:
@@ -404,7 +416,14 @@ async def process_next_file_task():
             
             if not image_paths:
                 print(f"[{FLOW_ID}] No images found in {zip_full_name}. Marking as complete.")
-                global_success = True
+                # Mark as processed and move to next
+                async with state_lock:
+                    state.processed_files.add(repo_file_full_path)
+                    state.current_index += 1
+                    state.current_file = None
+                    state.current_file_progress = "0/0"
+                    state.status = "Idle"
+                    await save_state_to_hf()
             else:
                 # Initialize progress tracker
                 progress_tracker = {
@@ -415,7 +434,7 @@ async def process_next_file_task():
                     state.current_file_progress = f"0/{len(image_paths)}"
                     await save_state_to_hf()
 
-                # Create and run captioning tasks
+                # Process images
                 semaphore = asyncio.Semaphore(len(servers))
                 async def limited_send_image_for_captioning(image_path, course_name, progress_tracker):
                     async with semaphore:
@@ -426,25 +445,14 @@ async def process_next_file_task():
                 all_captions = [r for r in results if r is not None]
                 
                 # Final progress report
-                if len(all_captions) == len(image_paths):
-                    print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: Successfully completed all {len(all_captions)} captions.")
-                    global_success = True
-                else:
-                    print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: Completed with partial result: {len(all_captions)}/{len(image_paths)} captions.")
-                    global_success = False
+                success_rate = len(all_captions) / len(image_paths)
+                print(f"[{FLOW_ID}] FINAL PROGRESS for {zip_full_name}: {len(all_captions)}/{len(image_paths)} captions ({success_rate:.1%})")
                 
-                # Upload results
+                # Upload results if we have any captions
                 if all_captions and 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}.")
-                        # If upload is successful, we mark the file as processed, regardless of partial success
-                        # The uploaded JSON will reflect the actual number of captions
-                        if global_success:
-                            print(f"[{FLOW_ID}] Fully processed and uploaded: {zip_full_name}")
-                        else:
-                            print(f"[{FLOW_ID}] Partially processed but uploaded: {zip_full_name}. Needs manual review.")
-                        
-                        # Mark as processed only if upload succeeded
+                        print(f"[{FLOW_ID}] Successfully uploaded captions for {zip_full_name}")
+                        # Mark as processed regardless of partial success
                         async with state_lock:
                             state.processed_files.add(repo_file_full_path)
                             state.current_index += 1
@@ -452,42 +460,37 @@ async def process_next_file_task():
                             state.current_file_progress = "0/0"
                             state.status = "Idle"
                             await save_state_to_hf()
-                            
                     else:
-                        print(f"[{FLOW_ID}] Failed to upload captions for {zip_full_name}. Will retry this file later.")
-                        # Do NOT increment index or mark as processed, so it will be retried
+                        print(f"[{FLOW_ID}] Failed to upload captions for {zip_full_name}. Will retry.")
+                        # Don't increment index, will retry this file
                         async with state_lock:
-                            state.status = f"Error uploading captions for {zip_full_name}. Retrying later."
+                            state.status = f"Upload failed for {zip_full_name}. Retrying later."
                             await save_state_to_hf()
-                        # Wait before retrying to avoid immediate re-attempt on a transient error
-                        await asyncio.sleep(60) 
-                        
+                        await asyncio.sleep(30)  # Wait before retry
                 else:
-                    print(f"[{FLOW_ID}] No captions generated or zip_full_name is missing. Skipping upload for {zip_full_name}. Will retry later.")
-                    # Do NOT increment index or mark as processed
+                    print(f"[{FLOW_ID}] No captions generated for {zip_full_name}. Will retry.")
+                    # Don't increment index, will retry this file
                     async with state_lock:
-                        state.status = f"No captions generated for {zip_full_name}. Retrying later."
+                        state.status = f"No captions for {zip_full_name}. Retrying later."
                         await save_state_to_hf()
-                    await asyncio.sleep(60)
+                    await asyncio.sleep(30)  # Wait before retry
             
         except Exception as e:
             error_message = str(e)
-            print(f"[{FLOW_ID}] Critical error in process_next_file_task for {repo_file_full_path}: {error_message}")
+            print(f"[{FLOW_ID}] Error processing {repo_file_full_path}: {error_message}")
             async with state_lock:
-                state.status = f"CRITICAL ERROR for {Path(repo_file_full_path).name}. Retrying later. Error: {error_message[:50]}..."
+                state.status = f"Error processing {Path(repo_file_full_path).name}: {error_message[:100]}..."
                 await save_state_to_hf()
-            # Wait before retrying
-            await asyncio.sleep(60)
+            await asyncio.sleep(30)  # Wait before retry
             
         finally:
             # Cleanup temporary files
             if extract_dir and extract_dir.exists():
-                print(f"[{FLOW_ID}] Cleaned up temporary directory {extract_dir}.")
-                shutil.rmtree(extract_dir, ignore_errors=True)
-                
-            # If the loop is still running, wait a short time before checking for the next file
-            if state.is_running:
-                await asyncio.sleep(5) 
+                try:
+                    shutil.rmtree(extract_dir, ignore_errors=True)
+                    print(f"[{FLOW_ID}] Cleaned up temporary directory {extract_dir}.")
+                except Exception as e:
+                    print(f"[{FLOW_ID}] Error cleaning up {extract_dir}: {e}")
 
 # --- FastAPI App and Endpoints ---
 
@@ -509,47 +512,24 @@ async def startup_event():
 
 async def initialize_after_startup():
     """Initialize app components after server has started"""
-    global state
-    
     try:
-        print(f"[{FLOW_ID}] Starting initialization...")
+        await load_state_from_hf()
+        await update_file_list()
         
-        # Load state with timeout
-        try:
-            print(f"[{FLOW_ID}] Loading state from HF...")
-            await asyncio.wait_for(load_state_from_hf(), timeout=30.0)
-            print(f"[{FLOW_ID}] State loaded successfully")
-        except asyncio.TimeoutError:
-            print("WARNING: Timeout loading state from HF, using default state")
-        except Exception as e:
-            print(f"WARNING: Error loading state from HF: {e}, using default state")
-        
-        # Update file list with timeout  
-        try:
-            print(f"[{FLOW_ID}] Updating file list...")
-            await asyncio.wait_for(update_file_list(), timeout=60.0)
-            print(f"[{FLOW_ID}] File list updated successfully")
-        except asyncio.TimeoutError:
-            print("WARNING: Timeout updating file list, using existing state")
-        except Exception as e:
-            print(f"WARNING: Error updating file list: {e}, using existing state")
+        # Always start the processing task, but it will check is_running flag
+        asyncio.create_task(process_next_file_task())
+        print(f"[{FLOW_ID}] Processing task created. Current running state: {state.is_running}")
         
-        # Start processing if needed
+        # Auto-start if we have files to process
         async with state_lock:
             if state.current_index < state.total_files:
                 state.is_running = True
-                asyncio.create_task(process_next_file_task())
-                print(f"[{FLOW_ID}] Started processing from index {state.current_index}/{state.total_files}")
-            else:
-                state.is_running = False
-                print(f"[{FLOW_ID}] All files processed. Starting in Idle mode.")
+                print(f"[{FLOW_ID}] Auto-starting processing from index {state.current_index}")
             
     except Exception as e:
-        print(f"[{FLOW_ID}] Critical error during initialization: {e}")
+        print(f"[{FLOW_ID}] Error during initialization: {e}")
         import traceback
         traceback.print_exc()
-        async with state_lock:
-            state.is_running = False
 
 
 @app.get("/", response_class=HTMLResponse)
@@ -634,21 +614,31 @@ async def control_processing(request: Request, background_tasks: BackgroundTasks
     
     async with state_lock:
         if action == "start":
-            if not state.is_running and state.current_index < state.total_files:
+            if not state.is_running:
+                # Reset state if we're at the end
+                if state.current_index >= state.total_files:
+                    state.current_index = 0
+                    state.status = "Reset to start and processing..."
+                
                 state.is_running = True
                 state.status = "Processing started."
+                
+                # Start the processing task
                 background_tasks.add_task(process_next_file_task)
                 await save_state_to_hf()
+                
+                print(f"[{FLOW_ID}] Processing manually started from index {state.current_index}")
                 return {"status": "success", "message": "Processing loop started."}
-            elif state.current_index >= state.total_files:
-                return {"status": "error", "message": "Cannot start. All files have been processed."}
             else:
                 return {"status": "info", "message": "Processing is already running."}
+                
         elif action == "stop":
             if state.is_running:
                 state.is_running = False
                 state.status = "Processing stopped by user."
                 await save_state_to_hf()
+                
+                print(f"[{FLOW_ID}] Processing manually stopped")
                 return {"status": "success", "message": "Processing loop stopped."}
             else:
                 return {"status": "info", "message": "Processing is already stopped."}