AnyStory / src /anystory.py
Junjie96's picture
Upload 46 files
9c18e52 verified
raw
history blame
2.52 kB
import json
import os
import time
import gradio as gr
import requests
from src.log import logger
from src.util import download_images
anystory_url = os.getenv("ANYSTORY_URL")
anystory_api_key = os.getenv("ANYSTORY_DS_API_KEY")
anystory_model = os.getenv("ANYSTORY_MODEL")
def call_anystory(image_urls, prompt):
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Bearer {anystory_api_key}",
"X-DashScope-Async": "enable",
"X-DashScope-DataInspection": "enable",
}
data = {
"model": anystory_model,
"input": {
"image_urls": image_urls,
"prompt": prompt
},
"parameters": {
},
}
res = requests.post(anystory_url, data=json.dumps(data), headers=headers)
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'https://poc-dashscope.aliyuncs.com/api/v1/tasks/{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.")
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)
logger.info(f"task_id: {task_id}: Generate done.")
else:
logger.error(f'Fail to create Generation task: {res.content}')
raise gr.Error("Fail to create Generation task.")
return img_data