diff --git a/app.py b/app.py index f9586bd..5bf8607 100644 --- a/app.py +++ b/app.py @@ -1,3 +1,4 @@ +import asyncio import base64 import hashlib import json @@ -129,15 +130,36 @@ def expo_clear_images(): os.remove(os.path.join(root, file)) @app.route('/expo_postprocess', methods=['POST']) -def expo_postprocess(): +async def expo_postprocess(): + global on_postprocessing if on_postprocessing: return jsonify({"status": "error", "message": "Postprocessing in progress"}), 503 - print("Postprocessing") + on_postprocessing = True - expo_postprocess_main() - expo_clear_images() + print("Postprocessing") + + # Wait until the directories have the same files or timeout + if not await wait_for_files_to_match(expo_openpose_dir, expo_raw_sd_dir): + print("Timeout reached, proceeding with postprocessing") + + await asyncio.to_thread(expo_postprocess_main) + await asyncio.to_thread(expo_clear_images) on_postprocessing = False + return jsonify({"status": "success", "message": "Postprocessing completed"}), 200 + + +async def wait_for_files_to_match(dir1: str, dir2: str, timeout: int = 60, interval: int = 1) -> bool: + start_time = asyncio.get_event_loop().time() + while asyncio.get_event_loop().time() - start_time < timeout: + files1 = set(os.listdir(dir1)) + files2 = set(os.listdir(dir2)) + if files1 == files2: + return True + await asyncio.sleep(interval) + return False + + @app.route('/gen_image', methods=['POST']) def gen_image(): if request.is_json: