from bytedance_ark_batch_inference.client import BatchInferenceClient
def generate_image_urls(tos_client, bucket_name, image_keys, style, save_bucket, save_objects):
for image_key, save_object in zip(image_keys, save_objects):
save_bucket=base64.b64encode(save_bucket.encode("utf-8")).decode("utf-8"),
save_object=base64.b64encode(save_object.encode("utf-8")).decode("utf-8"),
image_url = tos_client.pre_signed_url(
http_method=tos.HttpMethodType.Http_Method_Get,
image_urls.append(image_url.signed_url)
def create_requests_json(image_urls):
for i, image_url in enumerate(image_urls, start=1):
"custom_id": f"request-{i}",
"content": "你是豆包,是由字节跳动开发的 AI 人工智能助手",
"text": "请用几个关键词解释一下这张图帮我解释一下这张图",
"image_url": {"url": image_url},
requests.append(request)
def create_tos_client(ak, sk, endpoint, region):
return tos.TosClientV2(ak, sk, endpoint, region)
def upload_json_to_tos(tos_client, input_bucket_name, object_key, json_data):
json_bytes = json.dumps(json_data, ensure_ascii=False).encode('utf-8')
tos_client.put_object(bucket=input_bucket_name, key=object_key, content=json_bytes)
async def run_batch_inference(input_bucket_name, output_bucket_name, object_key, output_key):
batch_inference_client = BatchInferenceClient()
response = await batch_inference_client.create_batch_inference_job(input_bucket_name, output_bucket_name, object_key, output_key)
response = await batch_inference_client.ListBatchInferenceJobs(['Running'])
TOS_ACCESS_KEY = "Your Access Key"
TOS_SECRET_KEY = "Your Secret Key"
endpoint = "Your Endpoint"
bucket_name = "Your Bucket"
image_keys = ["image1.jpeg", "image2.jpeg", "image3.jpeg"]
style = "image/resize,h_100/format,jpg"
save_bucket = "Save Bucket"
save_objects = ["temp1.jpeg", "temp2.jpeg", "temp3.jpeg"]
input_key = "input.jsonl"
output_key = "output.jsonl"
tos_client = create_tos_client(TOS_ACCESS_KEY, TOS_SECRET_KEY, endpoint, region)
image_urls = generate_image_urls(tos_client, bucket_name, image_keys, style, save_bucket, save_objects)
requests_data = create_requests_json(image_urls)
upload_json_to_tos(tos_client, bucket_name, input_key, requests_data)
await run_batch_inference(bucket_name, bucket_name, input_key, output_key)
if __name__ == '__main__':