yatman
/
Image_Gen_Server
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix edge cases

This commit is contained in:
zaqxs123456 2024-10-25 11:43:28 +08:00
parent f5f947f6c2
commit 56d930a854
2 changed files with 20 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
app.py
View File

@ -25,8 +25,6 @@ info = json.load(open('info.json'))
comfyui_address = info['comfyui_address'] comfyui_address = info['comfyui_address']
expo_raw_sd_dir = info['expo_raw_sd_dir'] expo_raw_sd_dir = info['expo_raw_sd_dir']
expo_openpose_dir = info['expo_openpose_dir'] expo_openpose_dir = info['expo_openpose_dir']
expo_postprocessed_dir = info['expo_postprocessed_dir']
expo_postprocess_temp_dir = info['expo_postprocess_temp_dir']
on_postprocessing = False on_postprocessing = False
on_testing = False on_testing = False
@ -142,6 +140,14 @@ async def expo_postprocess():
if not await wait_for_files_to_match(expo_openpose_dir, expo_raw_sd_dir): if not await wait_for_files_to_match(expo_openpose_dir, expo_raw_sd_dir):
print("Timeout reached, proceeding with postprocessing") print("Timeout reached, proceeding with postprocessing")
# Check if directories exist and are not empty
if not os.path.exists(expo_openpose_dir) or not os.listdir(expo_openpose_dir):
on_postprocessing = False
return jsonify({"status": "error", "message": "No images to process in expo_openpose_dir"}), 404
if not os.path.exists(expo_raw_sd_dir) or not os.listdir(expo_raw_sd_dir):
on_postprocessing = False
return jsonify({"status": "error", "message": "No images to process in expo_raw_sd_dir"}), 404
await asyncio.to_thread(expo_postprocess_main) await asyncio.to_thread(expo_postprocess_main)
await asyncio.to_thread(expo_clear_images) await asyncio.to_thread(expo_clear_images)
on_postprocessing = False on_postprocessing = False
@ -152,14 +158,23 @@ async def expo_postprocess():
async def wait_for_files_to_match(dir1: str, dir2: str, timeout: int = 60, interval: int = 1) -> bool: 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() start_time = asyncio.get_event_loop().time()
while asyncio.get_event_loop().time() - start_time < timeout: while asyncio.get_event_loop().time() - start_time < timeout:
files1 = set(os.listdir(dir1)) files1 = get_all_files(dir1)
files2 = set(os.listdir(dir2)) files2 = get_all_files(dir2)
if files1 == files2: if files1 == files2:
return True return True
await asyncio.sleep(interval) await asyncio.sleep(interval)
return False return False
def get_all_files(directory):
all_files = set()
for root, _, files in os.walk(directory):
for file in files:
# Store the relative path of the file
relative_path = os.path.relpath(os.path.join(root, file), directory)
all_files.add(relative_path)
return all_files
@app.route('/gen_image', methods=['POST']) @app.route('/gen_image', methods=['POST'])
def gen_image(): def gen_image():
if request.is_json: if request.is_json:

1
postprocessing.py
View File

@ -191,6 +191,7 @@ def current_session():
def expo_postprocess_main(): def expo_postprocess_main():
print("Postprocessing") print("Postprocessing")
os.makedirs(expo_postprocessed_dir, exist_ok=True)
os.makedirs(expo_postprocess_temp_dir, exist_ok=True) os.makedirs(expo_postprocess_temp_dir, exist_ok=True)
left_fencer_raw_image_dir = os.path.join(expo_raw_sd_dir, 'left_fencer') left_fencer_raw_image_dir = os.path.join(expo_raw_sd_dir, 'left_fencer')