import base64 import hashlib import json import random import uuid import cv2 from flask import Flask, request, jsonify import sys import os from PIL import Image import io import numpy as np import websocket import openpose_gen as opg from comfy_socket import get_images from postprocessing import expo_shuffle_image_steps, expo_add_to_background_image sys.path.append('./') app = Flask(__name__) info = json.load(open('info.json')) comfyui_address = info['comfyui_address'] expo_raw_sd_dir = info['expo_raw_sd_dir'] expo_openpose_dir = info['expo_openpose_dir'] expo_postprocessed_dir = info['expo_postprocessed_dir'] expo_postprocess_temp_dir = info['expo_postprocess_temp_dir'] @app.route('/expo_fencing_pose', methods=['POST']) def expo_fencing_pose(): if request.is_json: data = request.get_json() coordinates = data['coordinates'] canvas_size = data['canvas_size'] batch = data['batch'] step = data['step'] if coordinates is None or canvas_size is None or 'batch' not in data or 'step' not in data: return jsonify({"status": "error", "message": "Missing data"}), 422 openpose_image_path = opg.expo_save_bodypose(canvas_size[0], canvas_size[1], coordinates, batch, step) print(openpose_image_path) expo_fencer_prompt(openpose_image_path, batch, step) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 def expo_fencer_prompt(openpose_image_path, batch, step): prompt = json.loads(open("./prompts/fencer_03.json", "r", encoding="utf-8").read()) openpose_image_name = opg.upload_image(openpose_image_path) opg.upload_image("./images/ref_black.png", "ref_black.png") print(openpose_image_name) prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["29"]["inputs"]['image'] = "ref_black.png" prompt["17"]["inputs"]['image'] = openpose_image_name client_id = hashlib.sha256(str(random.getrandbits(256)).encode('utf-8')).hexdigest() ws = websocket.WebSocket() ws.connect("ws://{}/ws?clientId={}".format(comfyui_address, client_id)) images = get_images(ws, prompt, client_id) for node_id in images: for idx, image_data in enumerate(images[node_id]): image = Image.open(io.BytesIO(image_data)) image_path = os.path.join(expo_raw_sd_dir, f"{batch}_{step}.png") image.save(image_path) def expo_clear_images(): for file in os.listdir(expo_openpose_dir): os.remove(os.path.join(expo_openpose_dir, file)) for file in os.listdir(expo_raw_sd_dir): os.remove(os.path.join(expo_raw_sd_dir, file)) @app.route('/expo_postprocess', methods=['POST']) def expo_postprocess(): print("Postprocessing") os.makedirs(expo_postprocess_temp_dir, exist_ok=True) shuffled_images_paths = expo_shuffle_image_steps() background_path = os.path.join(expo_postprocess_temp_dir, 'background.png') if not os.path.exists(background_path): background = np.zeros((1000, 1500, 3), dtype=np.uint8) cv2.imwrite(background_path, background) expo_add_to_background_image(background_path, shuffled_images_paths[0][0], 0, 0) cv2.imwrite(os.path.join(expo_postprocessed_dir, 'postprocessed.png'), background) # expo_clear_images() @app.route('/gen_image', methods=['POST']) def gen_image(): if request.is_json: data = request.get_json() coordinates = data['coordinates'] canvas_size = data['canvas_size'] pid = data['pid'] if not coordinates or not canvas_size: return jsonify({"status": "error", "message": "Missing data"}), 422 openpose_image_path = opg.save_bodypose(canvas_size[0], canvas_size[1], coordinates, pid) # gen_fencer_prompt(openpose_image_path, pid, comfyui_address) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 @app.route('/gen_group_pic', methods=['POST']) def gen_group_pic(): if request.is_json: data = request.get_json() coordinates_list = data['coordinates_list'] canvas_size = data['canvas_size'] pid = data['pid'] base_image = base64.b64decode(data['base_image']) if not coordinates_list or not canvas_size or not base_image or not pid: return jsonify({"status": "error", "message": "Missing data"}), 422 for i in range(len(coordinates_list)): coordinates_list[i] = coordinates_list[i]['coordinates'] openpose_image_path = opg.save_bodypose_mulit(canvas_size[0], canvas_size[1], coordinates_list, pid) gen_group_pic_prompt(openpose_image_path, base_image, pid, comfyui_address) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 def gen_fencer_prompt(openpose_image_path, pid, comfyUI_address): with open("./prompts/fencerAPI.json", "r") as f: prompt_json = f.read() prompt = json.loads(prompt_json) openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 20, pid, comfyUI_address) opg.upload_image("./images/ref_black.png", "ref_black.png") prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["29"]["inputs"]['image'] = "./images/ref_black.png" prompt["17"]["inputs"]['image'] = openpose_image_name opg.queue_prompt(prompt, comfyUI_address) def gen_group_pic_prompt(openpose_image_path, base_image, pid, comfyUI_address): with open("./prompts/group_pic.json", "r") as f: prompt_json = f.read() prompt = json.loads(prompt_json) openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 30, pid, comfyUI_address) base_image_name = opg.upload_image_circular_queue(base_image, 30, pid, comfyUI_address) prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["10"]["inputs"]['image'] = openpose_image_name prompt["14"]["inputs"]['image'] = base_image_name opg.queue_prompt(prompt, comfyUI_address) if __name__ == '__main__': expo_postprocess() # app.run(debug=True)