|
|
import os |
|
|
import json |
|
|
import time |
|
|
import uuid |
|
|
import hashlib |
|
|
import subprocess |
|
|
import huggingface_hub as hh |
|
|
from flask import Flask, request, send_file |
|
|
|
|
|
MAX_SECONDS = 100 |
|
|
server_me = 'https://zaisheng-testapi.hf.space' |
|
|
|
|
|
def get_key_sha256(key): |
|
|
sha256 = hashlib.sha256() |
|
|
sha256.update(key.encode('utf-8')) |
|
|
return sha256.hexdigest() |
|
|
|
|
|
match_sha256 = '686467fc40941a4af9b605365077c58ac06413870ec1679f41660047f716b7a6' |
|
|
|
|
|
prompt_folder = '../prompts' |
|
|
video_folder = '../videos' |
|
|
for f in [prompt_folder, video_folder]: |
|
|
os.makedirs(f, exist_ok=True) |
|
|
|
|
|
app = Flask(__name__, static_folder=video_folder, static_url_path='/generaged') |
|
|
|
|
|
|
|
|
@app.route("/v1/videos/image2video", methods = ['POST']) |
|
|
def image2video(): |
|
|
prompt_file_name = str(time.time()) |
|
|
h = request.headers |
|
|
key = h['Key'] |
|
|
if not get_key_sha256(key) == match_sha256: |
|
|
return 'err' |
|
|
d = request.json |
|
|
prompt = d['prompt'] |
|
|
open(os.path.join(prompt_folder, prompt_file_name), 'w').write(prompt) |
|
|
request_id = str(uuid.uuid4()) |
|
|
video_file = os.path.join(video_folder, prompt_file_name+'.mp4') |
|
|
for _ in range(MAX_SECONDS): |
|
|
time.sleep(1) |
|
|
if os.path.exists(video_file): |
|
|
updated_at = str(time.time()) |
|
|
ret = { |
|
|
"code": 0, |
|
|
"message": "The video is generated", |
|
|
"request_id": request_id, |
|
|
"data":{ |
|
|
"video_url":server_me+'/generaged/'+prompt_file_name+'.mp4', |
|
|
"task_id": request_id, |
|
|
"task_info":{ |
|
|
"external_task_id": request_id |
|
|
}, |
|
|
"task_status": "success", |
|
|
"created_at": prompt_file_name, |
|
|
"updated_at": updated_at |
|
|
} |
|
|
} |
|
|
return json.dumps(ret) |
|
|
updated_at = str(time.time()) |
|
|
ret = { |
|
|
"code": 1234, |
|
|
"message": "time out", |
|
|
"request_id": request_id, |
|
|
"data":{ |
|
|
"task_id": request_id, |
|
|
"task_info":{ |
|
|
"external_task_id": request_id |
|
|
}, |
|
|
"task_status": "fail", |
|
|
"created_at": prompt_file_name, |
|
|
"updated_at": updated_at |
|
|
} |
|
|
} |
|
|
return json.dumps(ret) |
|
|
|
|
|
|
|
|
@app.route("/getprompt", methods = ['GET']) |
|
|
def getprompt(): |
|
|
h = request.headers |
|
|
key = h['Key'] |
|
|
if not get_key_sha256(key) == match_sha256: |
|
|
return 'err' |
|
|
fs = sorted(os.listdir(prompt_folder)) |
|
|
fs = [f for f in fs if not f.startswith('.')] |
|
|
if len(fs) > 0: |
|
|
f = fs[0] |
|
|
ff = os.path.join(prompt_folder, f) |
|
|
txt = open(ff).read() |
|
|
d = {'file': f, 'prompt': txt} |
|
|
s = f'{f}||||{txt}' |
|
|
return s |
|
|
return "" |
|
|
|
|
|
@app.route("/putvideo", methods = ['POST']) |
|
|
def putvideo(): |
|
|
h = request.headers |
|
|
key = h['Key'] |
|
|
file = h['File'] |
|
|
repo_id_video = h['RV'] |
|
|
repo_id_prompt = h['RP'] |
|
|
tt = h['tt'] |
|
|
if not get_key_sha256(key) == match_sha256: |
|
|
return 'err' |
|
|
f = request.files['file'] |
|
|
local_file_video = os.path.join(video_folder, file+'.mp4') |
|
|
local_file_prompt = os.path.join(prompt_folder, file) |
|
|
f.save(local_file_video) |
|
|
hh.upload_file(path_or_fileobj=local_file_video, path_in_repo=file+'.mp4', repo_id=repo_id_video, repo_type='dataset', token=tt) |
|
|
hh.upload_file(path_or_fileobj=local_file_prompt, path_in_repo=file, repo_id=repo_id_prompt, repo_type='dataset', token=tt) |
|
|
os.remove(local_file_prompt) |
|
|
return "ok" |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
app.run(host="0.0.0.0", port=7860, debug=False) |
|
|
|
|
|
|
