import os import json import time import asyncio import aiohttp import zipfile from typing import Dict, List, Set, Optional from urllib.parse import quote from datetime import datetime from pathlib import Path import io from fastapi import FastAPI, BackgroundTasks, HTTPException, status from pydantic import BaseModel, Field from huggingface_hub import HfApi, hf_hub_download import uvicorn # --- Configuration --- # Flow Server ID and Port will be set via environment variables for easy deployment FLOW_ID = os.getenv("FLOW_ID", "flow_default") FLOW_PORT = int(os.getenv("FLOW_PORT", 8001)) # Default to 8001 for flow1 # Manager Server Configuration MANAGER_URL = os.getenv("MANAGER_URL", "https://fred808-fcord.hf.space") MANAGER_COMPLETE_TASK_URL = f"{MANAGER_URL}/task/complete" # Hugging Face Configuration HF_TOKEN = os.getenv("HF_TOKEN", "") # User provided token HF_DATASET_ID = os.getenv("HF_DATASET_ID", "Fred808/BG3") HF_OUTPUT_DATASET_ID = os.getenv("HF_OUTPUT_DATASET_ID", "fred808/helium") # Target dataset for captions # Using the full list from the user's original code for actual deployment CAPTION_SERVERS = [ "https://fred808-pil-4-1.hf.space/analyze", "https://fred808-pil-4-2.hf.space/analyze", "https://fred808-pil-4-3.hf.space/analyze", "https://fred1012-fred1012-gw0j2h.hf.space/analyze", "https://fred1012-fred1012-wqs6c2.hf.space/analyze", "https://fred1012-fred1012-oncray.hf.space/analyze", "https://fred1012-fred1012-4goge7.hf.space/analyze", "https://fred1012-fred1012-z0eh7m.hf.space/analyze", "https://fred1012-fred1012-u95rte.hf.space/analyze", "https://fred1012-fred1012-igje22.hf.space/analyze", "https://fred1012-fred1012-ibkuf8.hf.space/analyze", "https://fred1012-fred1012-nwqthy.hf.space/analyze", "https://fred1012-fred1012-4ldqj4.hf.space/analyze", "https://fred1012-fred1012-pivlzg.hf.space/analyze", "https://fred1012-fred1012-ptlc5u.hf.space/analyze", "https://fred1012-fred1012-u7lh57.hf.space/analyze", "https://fred1012-fred1012-q8djv1.hf.space/analyze", "https://fredalone-fredalone-ozugrp.hf.space/analyze", "https://fredalone-fredalone-9brxj2.hf.space/analyze", "https://fredalone-fredalone-p8vq9a.hf.space/analyze", "https://fredalone-fredalone-vbli2y.hf.space/analyze", "https://fredalone-fredalone-uggger.hf.space/analyze", "https://fredalone-fredalone-nmi7e8.hf.space/analyze", "https://fredalone-fredalone-d1f26d.hf.space/analyze", "https://fredalone-fredalone-461jp2.hf.space/analyze", "https://fredalone-fredalone-3enfg4.hf.space/analyze", "https://fredalone-fredalone-dqdbpv.hf.space/analyze", "https://fredalone-fredalone-ivtjua.hf.space/analyze", "https://fredalone-fredalone-6bezt2.hf.space/analyze", "https://fredalone-fredalone-e0wfnk.hf.space/analyze", "https://fredalone-fredalone-zu2t7j.hf.space/analyze", "https://fredalone-fredalone-dqtv1o.hf.space/analyze", "https://fredalone-fredalone-wclyog.hf.space/analyze", "https://fredalone-fredalone-t27vig.hf.space/analyze", "https://fredalone-fredalone-gahbxh.hf.space/analyze", "https://fredalone-fredalone-kw2po4.hf.space/analyze", "https://fredalone-fredalone-8h285h.hf.space/analyze" ] MODEL_TYPE = "Florence-2-large" # Temporary storage for images TEMP_DIR = Path(f"temp_images_{FLOW_ID}") TEMP_DIR.mkdir(exist_ok=True) # --- Models --- class ProcessCourseRequest(BaseModel): course_name: Optional[str] = None class CaptionServer: def __init__(self, url): self.url = url self.busy = False self.total_processed = 0 self.total_time = 0 self.model = MODEL_TYPE @property def fps(self): return self.total_processed / self.total_time if self.total_time > 0 else 0 # Global state for caption servers servers = [CaptionServer(url) for url in CAPTION_SERVERS] server_index = 0 # --- Core Processing Functions --- async def get_available_server(timeout: float = 300.0) -> CaptionServer: """Round-robin selection of an available caption server.""" global server_index start_time = time.time() while True: # Round-robin check for an available server for _ in range(len(servers)): server = servers[server_index] server_index = (server_index + 1) % len(servers) if not server.busy: return server # If all servers are busy, wait for a short period and check again await asyncio.sleep(0.5) # Check if timeout has been reached if time.time() - start_time > timeout: raise TimeoutError(f"Timeout ({timeout}s) waiting for an available caption server.") async def send_image_for_captioning(image_path: Path, course_name: str, progress_tracker: Dict) -> Optional[Dict]: """Sends a single image to a caption server for processing.""" # This function now handles server selection and retries internally MAX_RETRIES = 3 for attempt in range(MAX_RETRIES): server = None try: # 1. Get an available server (will wait if all are busy, with a timeout) server = await get_available_server() server.busy = True start_time = time.time() # Print a less verbose message only on the first attempt if attempt == 0: print(f"[{FLOW_ID}] Starting attempt on {image_path.name}...") # 2. Prepare request data form_data = aiohttp.FormData() form_data.add_field('file', image_path.open('rb'), filename=image_path.name, content_type='image/jpeg') form_data.add_field('model_choice', MODEL_TYPE) # 3. Send request async with aiohttp.ClientSession() as session: # Increased timeout to 10 minutes (600s) as requested by user's problem description async with session.post(server.url, data=form_data, timeout=600) as resp: if resp.status == 200: result = await resp.json() caption = result.get("caption") if caption: # Update progress counter progress_tracker['completed'] += 1 if progress_tracker['completed'] % 50 == 0: print(f"[{FLOW_ID}] PROGRESS: {progress_tracker['completed']}/{progress_tracker['total']} captions completed.") # Log success only if it's not a progress report interval if progress_tracker['completed'] % 50 != 0: print(f"[{FLOW_ID}] Success: {image_path.name} captioned by {server.url}") return { "course": course_name, "image_path": image_path.name, "caption": caption, "timestamp": datetime.now().isoformat() } else: print(f"[{FLOW_ID}] Server {server.url} returned success but no caption for {image_path.name}. Retrying...") continue # Retry with a different server else: error_text = await resp.text() print(f"[{FLOW_ID}] Error from server {server.url} for {image_path.name}: {resp.status} - {error_text}. Retrying...") continue # Retry with a different server except (aiohttp.ClientError, asyncio.TimeoutError, TimeoutError) as e: print(f"[{FLOW_ID}] Connection/Timeout error for {image_path.name} on {server.url if server else 'unknown server'}: {e}. Retrying...") continue # Retry with a different server except Exception as e: print(f"[{FLOW_ID}] Unexpected error during captioning for {image_path.name}: {e}. Retrying...") continue # Retry with a different server finally: if server: end_time = time.time() server.busy = False server.total_processed += 1 server.total_time += (end_time - start_time) print(f"[{FLOW_ID}] FAILED after {MAX_RETRIES} attempts for {image_path.name}.") return None async def download_and_extract_zip(course_name: str, processed_files: Set[str]) -> Optional[tuple[Path, str, str]]: """Downloads the zip file for the course and extracts its contents.""" print(f"[{FLOW_ID}] Looking for files starting with '{course_name}' in frames/ directory...") try: api = HfApi(token=HF_TOKEN) # List all files in the frames directory repo_files = api.list_repo_files( repo_id=HF_DATASET_ID, repo_type="dataset" ) # Find zip files that start with the course name matching_files = [ f for f in repo_files if f.startswith(f"frames/{course_name}") and f.endswith('.zip') ] if not matching_files: print(f"[{FLOW_ID}] No zip files found starting with '{course_name}' in frames/ directory.") return None, None # Filter out already processed files and select the first one unprocessed_files = [f for f in matching_files if f not in processed_files] if not unprocessed_files: print(f"[{FLOW_ID}] No new zip files found for '{course_name}'.") return None, None, None repo_file_full_path = unprocessed_files[0] # e.g., frames/DAREEFSA_full_name.zip # Extract the full file name from the path (e.g., DAREEFSA_full_name.zip) zip_full_name = Path(repo_file_full_path).name print(f"[{FLOW_ID}] Found new matching file: {repo_file_full_path}. Full name: {zip_full_name}") # Use hf_hub_download to get the file path zip_path = hf_hub_download( repo_id=HF_DATASET_ID, filename=repo_file_full_path, # Use the full path in the repo repo_type="dataset", token=HF_TOKEN, ) print(f"[{FLOW_ID}] Downloaded to {zip_path}. Extracting...") # Create a temporary directory for extraction extract_dir = TEMP_DIR / course_name extract_dir.mkdir(exist_ok=True) with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(extract_dir) print(f"[{FLOW_ID}] Extraction complete to {extract_dir}.") # Return the extraction directory, the full zip file name, and the repo path return extract_dir, zip_full_name, repo_file_full_path except Exception as e: print(f"[{FLOW_ID}] Error downloading or extracting zip for {course_name}: {e}") return None, None, None async def upload_captions_to_hf(zip_full_name: str, captions: List[Dict]) -> bool: """Uploads the final captions JSON file to the output dataset. The user requested the output JSON file to be named after the full zip file name. """ # Use the full zip name, replacing the extension with .json caption_filename = Path(zip_full_name).with_suffix('.json').name try: print(f"[{FLOW_ID}] Uploading {len(captions)} captions for {zip_full_name} as {caption_filename} to {HF_OUTPUT_DATASET_ID}...") # Create JSON content in memory json_content = json.dumps(captions, indent=2, ensure_ascii=False).encode('utf-8') api = HfApi(token=HF_TOKEN) api.upload_file( path_or_fileobj=io.BytesIO(json_content), path_in_repo=caption_filename, repo_id=HF_OUTPUT_DATASET_ID, repo_type="dataset", commit_message=f"[{FLOW_ID}] Captions for {zip_full_name}" ) print(f"[{FLOW_ID}] Successfully uploaded captions for {zip_full_name}.") return True except Exception as e: print(f"[{FLOW_ID}] Error uploading captions for {zip_full_name}: {e}") return False async def process_course_task(course_name: str): """Main task to process a single course, looping until all files are processed.""" print(f"[{FLOW_ID}] Starting continuous processing for course: {course_name}") processed_files = set() all_processed_files_log = [] global_success = True # Loop to continuously check for new files matching the course_name prefix while True: extract_dir = None zip_full_name = None repo_file_full_path = None try: # download_and_extract_zip now returns a tuple: (extract_dir, zip_full_name, repo_file_full_path) download_result = await download_and_extract_zip(course_name, processed_files) if download_result is None or download_result[0] is None: # No new files found, or an error occurred during search/download if download_result is not None and download_result[0] is None and download_result[1] is None: print(f"[{FLOW_ID}] No new files found for {course_name}. Exiting loop.") break else: # An error occurred during search/download raise Exception("Failed to download or extract zip file.") extract_dir, zip_full_name, repo_file_full_path = download_result # Add the file to the processed set immediately to avoid re-processing in the next loop processed_files.add(repo_file_full_path) all_processed_files_log.append(repo_file_full_path) # --- Start Processing the single file --- # FIX: Use recursive glob to find images in subdirectories image_paths = [p for p in extract_dir.glob("**/*") if p.is_file() and p.suffix.lower() in ['.jpg', '.jpeg', '.png']] print(f"[{FLOW_ID}] Found {len(image_paths)} images to process in {zip_full_name}.") current_file_success = False if not image_paths: print(f"[{FLOW_ID}] No images found in {zip_full_name}. Marking as complete.") current_file_success = True else: # Initialize progress tracker progress_tracker = { 'total': len(image_paths), 'completed': 0 } print(f"[{FLOW_ID}] Starting captioning for {progress_tracker['total']} images in {zip_full_name}...") # Create a semaphore to limit concurrent tasks to the number of available servers semaphore = asyncio.Semaphore(len(servers)) async def limited_send_image_for_captioning(image_path, course_name, progress_tracker): async with semaphore: return await send_image_for_captioning(image_path, course_name, progress_tracker) # Create a list of tasks for parallel captioning caption_tasks = [] for image_path in image_paths: caption_tasks.append(limited_send_image_for_captioning(image_path, course_name, progress_tracker)) # Run all captioning tasks concurrently results = await asyncio.gather(*caption_tasks) # Filter out failed results all_captions = [r for r in results if r is not None] # 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.") current_file_success = False # Upload results if all_captions and zip_full_name: # Use the full zip file name for the upload as requested 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}.") # If partial success, we still upload, but the overall task is marked as failure if any file failed if not current_file_success: global_success = False else: print(f"[{FLOW_ID}] Failed to upload captions for {zip_full_name}.") current_file_success = False global_success = False else: print(f"[{FLOW_ID}] No captions generated or zip_full_name is missing. Skipping upload for {zip_full_name}.") current_file_success = False global_success = False # --- End Processing the single file --- except Exception as e: error_message = str(e) print(f"[{FLOW_ID}] Critical error in process_course_task for {course_name}: {error_message}") global_success = False finally: # Cleanup temporary files for the current file if extract_dir and extract_dir.exists(): print(f"[{FLOW_ID}] Cleaned up temporary directory {extract_dir}.") import shutil shutil.rmtree(extract_dir, ignore_errors=True) # If an unrecoverable error occurred (e.g., during search/download), break the loop if download_result is None and extract_dir is None: break # --- Final Report after the loop is complete --- print(f"[{FLOW_ID}] All processing loops complete for {course_name}.") print(f"[{FLOW_ID}] Total files processed: {len(all_processed_files_log)}") print(f"[{FLOW_ID}] List of processed files: {all_processed_files_log}") # Report completion to manager final_error_message = error_message if not global_success else None # Assuming report_completion exists and is an async function # await report_completion(course_name, global_success, final_error_message) return global_success async def report_completion(course_name: str, success: bool, error_message: Optional[str] = None): """Reports the task result back to the Manager Server.""" print(f"[{FLOW_ID}] Reporting completion for {course_name} (Success: {success})...") payload = { "flow_id": FLOW_ID, "course_name": course_name, "success": success, "error_message": error_message } try: async with aiohttp.ClientSession() as session: async with session.post(MANAGER_COMPLETE_TASK_URL, json=payload) as resp: if resp.status != 200: print(f"[{FLOW_ID}] ERROR: Manager reported non-200 status: {resp.status} - {await resp.text()}") else: print(f"[{FLOW_ID}] Successfully reported completion to Manager.") except aiohttp.ClientError as e: print(f"[{FLOW_ID}] CRITICAL ERROR: Could not connect to Manager at {MANAGER_COMPLETE_TASK_URL}. Task completion not reported. Error: {e}") except Exception as e: print(f"[{FLOW_ID}] Unexpected error during reporting: {e}") # --- FastAPI App and Endpoints --- app = FastAPI( title=f"Flow Server {FLOW_ID} API", description="Fetches, extracts, and captions images for a given course.", version="1.0.0" ) @app.on_event("startup") async def startup_event(): print(f"Flow Server {FLOW_ID} started on port {FLOW_PORT}. Manager URL: {MANAGER_URL}") @app.get("/") async def root(): return { "flow_id": FLOW_ID, "status": "ready", "manager_url": MANAGER_URL, "total_servers": len(servers), "busy_servers": sum(1 for s in servers if s.busy), } @app.post("/process_course") async def process_course(request: ProcessCourseRequest, background_tasks: BackgroundTasks): """ Receives a course name from the Manager and starts processing in the background. """ course_name = request.course_name if not course_name: print(f"[{FLOW_ID}] Received empty course name. Stopping processing loop.") return {"status": "stopped", "message": "No more courses to process."} print(f"[{FLOW_ID}] Received course: {course_name}. Starting background task.") # Start the heavy processing in a background task so the API call returns immediately background_tasks.add_task(process_course_task, course_name) return {"status": "processing", "course_name": course_name, "message": "Processing started in background."} if __name__ == "__main__": # Note: When running in the sandbox, we need to use 0.0.0.0 to expose the port. uvicorn.run(app, host="0.0.0.0", port=FLOW_PORT)