Spaces:
Running
Running
File size: 3,048 Bytes
dbac7c5 6475329 2538f2a dbac7c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import json
import os
import time
import gradio as gr
import requests
from src.log import logger
from src.util import download_images
def call_generation(data):
url_task = os.getenv("URL_TASK")
api_key = os.getenv("API_KEY_GENERATION")
model_id = os.getenv("MODEL_ID")
url_query = os.getenv("URL_QUERY")
batch_size = 4
repeat_times = 1
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Bearer {api_key}",
"X-DashScope-Async": "enable",
}
data["model"] = model_id
data["parameters"]["n"] = batch_size
all_res_ = []
for i in range(repeat_times):
if data["parameters"]["seed"] != -1:
data["parameters"]["seed"] = data["parameters"]["seed"] * (i+1)
res_ = requests.post(url_task, data=json.dumps(data), headers=headers)
all_res_.append(res_)
all_image_data = []
for res_ in all_res_:
respose_code = res_.status_code
if 200 == respose_code:
res = json.loads(res_.content.decode())
task_id = res['output']['task_id']
logger.info(f"task_id: {task_id}: Create request success. Params: {data}")
# Async query
is_running = True
while is_running:
res_ = requests.post(f'{url_query}/{task_id}', headers=headers)
respose_code = res_.status_code
if 200 == respose_code:
res = json.loads(res_.content.decode())
if "SUCCEEDED" == res['output']['task_status']:
logger.info(f"task_id: {task_id}: Generation task query success.")
results = res['output']['results']
img_urls = [x['url'] for x in results]
logger.info(f"task_id: {task_id}: {res}")
break
elif "FAILED" != res['output']['task_status']:
logger.debug(f"task_id: {task_id}: query result...")
time.sleep(1)
else:
raise gr.Error(
"Fail to get results from Generation task. Make sure all the ID images have a clear face. If it still doesn't work, you can contact us or open an issue.")
else:
logger.error(f'task_id: {task_id}: Fail to query task result: {res_.content}')
raise gr.Error("Fail to query task result.")
logger.info(f"task_id: {task_id}: download generated images.")
img_data = download_images(img_urls, batch_size)
logger.info(f"task_id: {task_id}: Generate done.")
all_image_data += img_data
else:
logger.error(f'Fail to create Generation task: {res_.content}')
raise gr.Error("Fail to create Generation task.")
if len(all_image_data) != repeat_times * batch_size:
raise gr.Error("Fail to Generation.")
return all_image_data
|