This commit is contained in:
iHeyTang
2025-08-25 14:42:35 +08:00
parent 467c8b2549
commit f41fd24316
17 changed files with 4353 additions and 943 deletions

View File

@@ -13,10 +13,8 @@ from torch import Tensor
def task_submit(uid, video, audio, heygem_url):
"""Submit a task to the API"""
task_submit_api = f'{heygem_url}/easy/submit'
result_json = {
'status': False, 'data': {}, 'msg': ''
}
task_submit_api = f"{heygem_url}/easy/submit"
result_json = {"status": False, "data": {}, "msg": ""}
try:
data = {
"code": uid,
@@ -24,104 +22,112 @@ def task_submit(uid, video, audio, heygem_url):
"audio_url": audio,
"chaofen": 1,
"watermark_switch": 0,
"pn": 1
"pn": 1,
}
loguru.logger.info(f'data={data}')
loguru.logger.info(f"data={data}")
with httpx.Client() as client:
resp = client.post(task_submit_api, json=data)
resp_dict = resp.json()
loguru.logger.info(f'submit data: {resp_dict}')
if resp_dict['code'] != 10000:
result_json['status'] = False
result_json['msg'] = result_json['msg']
loguru.logger.info(f"submit data: {resp_dict}")
if resp_dict["code"] != 10000:
result_json["status"] = False
result_json["msg"] = result_json["msg"]
else:
result_json['status'] = True
result_json['data'] = uid
result_json['msg'] = '任务提交成功'
result_json["status"] = True
result_json["data"] = uid
result_json["msg"] = "任务提交成功"
except Exception as e:
loguru.logger.info(f'submit task fail case by:{str(e)}')
loguru.logger.info(f"submit task fail case by:{str(e)}")
raise RuntimeError(str(e))
return result_json
def query_task_progress(heygem_url, heygem_temp_path, task_id: str, interval: int = 10, timeout: int = 60 * 15):
def query_task_progress(
heygem_url,
heygem_temp_path,
task_id: str,
interval: int = 10,
timeout: int = 60 * 15,
):
"""Query task progress and wait for completion"""
result_json = {'status': False, 'data': {}, 'msg': ''}
result_json = {"status": False, "data": {}, "msg": ""}
def query_result(t_id: str):
tmp_dict = {'status': True, 'data': dict(), 'msg': ''}
tmp_dict = {"status": True, "data": dict(), "msg": ""}
try:
query_task_url = f'{heygem_url}/easy/query'
params = {
'code': t_id
}
query_task_url = f"{heygem_url}/easy/query"
params = {"code": t_id}
with httpx.Client() as client:
resp = client.get(query_task_url, params=params)
resp_dict = resp.json()
status_code = resp_dict['code']
status_code = resp_dict["code"]
if status_code in (9999, 10002, 10003, 10001):
tmp_dict['status'] = False
tmp_dict['msg'] = resp_dict['msg']
tmp_dict["status"] = False
tmp_dict["msg"] = resp_dict["msg"]
elif status_code == 10000:
loguru.logger.info(f'query task data: {json.dumps(resp_dict)}')
status_code = resp_dict['data'].get('status', 1)
loguru.logger.info(f"query task data: {json.dumps(resp_dict)}")
status_code = resp_dict["data"].get("status", 1)
if status_code == 3:
tmp_dict['status'] = False
tmp_dict['msg'] = resp_dict['data']['msg']
tmp_dict["status"] = False
tmp_dict["msg"] = resp_dict["data"]["msg"]
else:
process = resp_dict['data'].get('progress', 20)
process = resp_dict["data"].get("progress", 20)
if status_code == 2:
process = 100
else:
process = process
result = resp_dict['data'].get('result', '')
tmp_dict['data'] = {'progress': process,
'path': result,
}
result = resp_dict["data"].get("result", "")
tmp_dict["data"] = {
"progress": process,
"path": result,
}
else:
pass
except Exception as e:
loguru.logger.info(f'query task fail case by:{str(e)}')
loguru.logger.info(f"query task fail case by:{str(e)}")
raise RuntimeError(str(e))
return tmp_dict
end = time.time() + timeout
while time.time() < end:
tmp_result = query_result(task_id)
if not tmp_result['status'] or tmp_result['data'].__eq__({}):
result_json['status'] = False
result_json['msg'] = tmp_result['msg']
if not tmp_result["status"] or tmp_result["data"].__eq__({}):
result_json["status"] = False
result_json["msg"] = tmp_result["msg"]
break
else:
process = tmp_result['data']['progress']
loguru.logger.info(f'query task progress :{process}')
if tmp_result['data']['progress'] < 100:
process = tmp_result["data"]["progress"]
loguru.logger.info(f"query task progress :{process}")
if tmp_result["data"]["progress"] < 100:
time.sleep(interval)
loguru.logger.info(f'wait next interval:{interval}')
loguru.logger.info(f"wait next interval:{interval}")
else:
p = tmp_result['data']['path']
p = p.replace('/', '').replace('\\', '')
result_json['data'] = "%s/%s"%(heygem_temp_path, p)
result_json['status'] = True
p = tmp_result["data"]["path"]
p = p.replace("/", "").replace("\\", "")
result_json["data"] = "%s/%s" % (heygem_temp_path, p)
result_json["status"] = True
return result_json
return result_json
def path_convert(path):
if ":" in path:
path = path.replace(os.sep,"/").split(":")
path = path.replace(os.sep, "/").split(":")
path[0] = path[0].lower()
path[1] = path[1][1:]
path = "/".join(["/mnt",*path])
path = "/".join(["/mnt", *path])
return path
def result_path_convert(result_path:str):
def result_path_convert(result_path: str):
if result_path.startswith("/"):
result_path = result_path.replace("/","\\")
result_path = result_path.replace("/", "\\")
result_path = r"\\wsl.localhost\Debian" + result_path
return result_path
class HeyGemF2F:
"""HeyGem 嘴型同步"""
@@ -132,8 +138,8 @@ class HeyGemF2F:
"video": ("IMAGE", {"forceInput": True}),
"audio": ("AUDIO", {"forceInput": True}),
"heygem_url": ("STRING", {"default": "http://127.0.0.1:8383"}),
"heygem_temp_path": ("STRING", {"default":"/code/data/temp"}),
"is_Windows": ("BOOLEAN", {"default": False})
"heygem_temp_path": ("STRING", {"default": "/code/data/temp"}),
"is_Windows": ("BOOLEAN", {"default": False}),
}
}
@@ -142,36 +148,49 @@ class HeyGemF2F:
FUNCTION = "f2f"
CATEGORY = "不忘科技-自定义节点🚩/视频/口型"
def f2f(self, video:Tensor, audio:dict, heygem_url:str, heygem_temp_path:str, is_Windows:bool):
def f2f(
self,
video: Tensor,
audio: dict,
heygem_url: str,
heygem_temp_path: str,
is_Windows: bool,
):
uid = str(uuid.uuid4())
video_path = os.path.join(os.path.dirname(__file__),"%s.mp4" % uid)
audio_path = os.path.join(os.path.dirname(__file__),"%s.wav" % uid)
video_path = os.path.join(os.path.dirname(__file__), "%s.mp4" % uid)
audio_path = os.path.join(os.path.dirname(__file__), "%s.wav" % uid)
try:
try:
torchvision.io.write_video(video_path, video.mul_(255).int(),25)
torchaudio.save(audio_path, audio["waveform"].squeeze(0), audio["sample_rate"], True)
torchvision.io.write_video(video_path, video.mul_(255).int(), 25)
torchaudio.save(
audio_path, audio["waveform"].squeeze(0), audio["sample_rate"], True
)
except:
traceback.print_exc()
raise RuntimeError("Save Temp File Error! ")
submit_result = task_submit(uid, path_convert(video_path), path_convert(audio_path), heygem_url)
if not submit_result['status']:
submit_result = task_submit(
uid, path_convert(video_path), path_convert(audio_path), heygem_url
)
if not submit_result["status"]:
return {
'status': False,
'data': {},
'msg': f"Task submission failed: {submit_result['msg']}"
"status": False,
"data": {},
"msg": f"Task submission failed: {submit_result['msg']}",
}
task_id = submit_result['data']
loguru.logger.info(f'Submitted task: {task_id}')
task_id = submit_result["data"]
loguru.logger.info(f"Submitted task: {task_id}")
# Query task progress
progress_result = query_task_progress(heygem_url, heygem_temp_path, task_id, interval=5)
progress_result = query_task_progress(
heygem_url, heygem_temp_path, task_id, interval=5
)
if not progress_result['status']:
if not progress_result["status"]:
raise RuntimeError(f"Task processing failed: {progress_result['msg']}")
# Return the file for download
file_path = progress_result['data']
file_path = progress_result["data"]
if is_Windows:
file_path = result_path_convert(file_path)
if os.path.exists(file_path):
@@ -183,11 +202,12 @@ class HeyGemF2F:
raise Exception(str(e))
finally:
try:
os.remove(os.path.join(os.path.dirname(__file__),"%s.mp4" % uid))
os.remove(os.path.join(os.path.dirname(__file__),"%s.wav" % uid))
os.remove(os.path.join(os.path.dirname(__file__), "%s.mp4" % uid))
os.remove(os.path.join(os.path.dirname(__file__), "%s.wav" % uid))
except:
pass
class HeyGemF2FFromFile:
"""HeyGem 嘴型同步 直接读取文件"""
@@ -198,8 +218,8 @@ class HeyGemF2FFromFile:
"video": ("STRING", {"forceInput": True}),
"audio": ("STRING", {"forceInput": True}),
"heygem_url": ("STRING", {"default": "http://127.0.0.1:8383"}),
"heygem_temp_path": ("STRING", {"default":"/code/data/temp"}),
"is_Windows": ("BOOLEAN", {"default": False})
"heygem_temp_path": ("STRING", {"default": "/code/data/temp"}),
"is_Windows": ("BOOLEAN", {"default": False}),
}
}
@@ -208,28 +228,35 @@ class HeyGemF2FFromFile:
FUNCTION = "f2f"
CATEGORY = "不忘科技-自定义节点🚩/视频/口型"
def f2f(self, video:str, audio:str, heygem_url:str, heygem_temp_path:str, is_Windows:bool):
def f2f(
self,
video: str,
audio: str,
heygem_url: str,
heygem_temp_path: str,
is_Windows: bool,
):
uid = str(uuid.uuid4())
try:
submit_result = task_submit(uid, video, audio, heygem_url)
if not submit_result['status']:
if not submit_result["status"]:
return {
'status': False,
'data': {},
'msg': f"Task submission failed: {submit_result['msg']}"
"status": False,
"data": {},
"msg": f"Task submission failed: {submit_result['msg']}",
}
task_id = submit_result['data']
loguru.logger.info(f'Submitted task: {task_id}')
task_id = submit_result["data"]
loguru.logger.info(f"Submitted task: {task_id}")
# Query task progress
progress_result = query_task_progress(heygem_url,heygem_temp_path,task_id)
progress_result = query_task_progress(heygem_url, heygem_temp_path, task_id)
if not progress_result['status']:
if not progress_result["status"]:
raise RuntimeError(f"Task processing failed: {progress_result['msg']}")
# Return the file for download
file_path = progress_result['data']
file_path = progress_result["data"]
if is_Windows:
file_path = result_path_convert(file_path)
if os.path.exists(file_path):
@@ -238,4 +265,4 @@ class HeyGemF2FFromFile:
raise FileNotFoundError(f"Output file not found at {file_path}")
except Exception as e:
loguru.logger.error(f"Error processing request: {str(e)}")
raise Exception(str(e))
raise Exception(str(e))