SIUBIU commited on
Commit
8c6d468
1 Parent(s): aab75c9

Upload kolorsvton1024.py

Browse files
Files changed (1) hide show
  1. kolorsvton1024.py +170 -0
kolorsvton1024.py ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ import jwt
3
+ import requests
4
+ import json
5
+ import base64
6
+ import os
7
+ def image_to_base64(image_path):
8
+ """
9
+ 将图片文件转换为Base64编码格式
10
+ :param image_path: 图片文件路径
11
+ :return: Base64编码字符串
12
+ """
13
+ with open(image_path, "rb") as image_file:
14
+ return base64.b64encode(image_file.read()).decode('utf-8')
15
+
16
+ ak = "c1748494bc5d42ed8db2b2d24ceb1a2b" # 填写access key
17
+ sk = "1cf0bbef768746f79d30ef13630b393b" # 填写secret key
18
+
19
+ def encode_jwt_token(ak, sk):
20
+ headers = {
21
+ "alg": "HS256",
22
+ "typ": "JWT"
23
+ }
24
+ payload = {
25
+ "iss": ak,
26
+ "exp": int(time.time()) + 1800, # 有效时间,此处示例代表当前时间+1800s(30min)
27
+ "nbf": int(time.time()) - 5 # 开始生效的时间,此处示例代表当前时间-5秒
28
+ }
29
+ token = jwt.encode(payload, sk, headers=headers)
30
+ return token
31
+
32
+
33
+
34
+ # ========== 配置参数 ==========
35
+
36
+ # API 请求的基础配置
37
+ api_token = encode_jwt_token(ak, sk)
38
+ print(api_token) # 打印生成的API_TOKEN
39
+ BASE_URL = "https://api.klingai.com/v1/images/kolors-virtual-try-on"
40
+
41
+ # 请求头
42
+ HEADERS = {
43
+ "Content-Type": "application/json",
44
+ "Authorization": f"Bearer {api_token}"
45
+ }
46
+
47
+ # 图片及回调配置
48
+ # HUMAN_IMAGE = image_to_base64(r"F:\codetest24\0903_lf\vton\people.jpg") # 替换为人物图片的Base64或URL
49
+ # CLOTH_IMAGE = image_to_base64(r"F:\codetest24\0903_lf\vton\hanfu.jpg") # 替换为服饰图片的Base64或URL
50
+ # CALLBACK_URL = "" # 可选,任务结果回调通知URL
51
+ SAVE_DIRECTORY = "downloads/kolor"
52
+ # ========== 步骤 1:创建虚拟试穿任务 ==========
53
+
54
+ def create_virtual_tryon_task(humen, cloth):
55
+ """
56
+ 创建虚拟试穿任务
57
+ """
58
+ # 请求体数据
59
+ data = {
60
+ "model_name": "kolors-virtual-try-on-v1",
61
+ "human_image": image_to_base64(humen),
62
+ "cloth_image": image_to_base64(cloth),
63
+ }
64
+
65
+ # 发起POST请求创建任务
66
+ response = requests.post(BASE_URL, headers=HEADERS, data=json.dumps(data))
67
+
68
+ # 处理响应
69
+ if response.status_code == 200:
70
+ result = response.json()
71
+ task_id = result['data']['task_id']
72
+ print(f"创建任务成功,任务ID: {task_id}")
73
+ return task_id
74
+ else:
75
+ print(f"创建任务失败: {response.status_code}, {response.text}")
76
+ return None
77
+
78
+ # ========== 步骤 2:查询单个虚拟试穿任务状态并保存图片 ==========
79
+
80
+ def query_virtual_tryon_task(task_id, i):
81
+ """
82
+ 根据任务ID查询任务状态,成功后保存生成的图片
83
+ """
84
+ # 请求URL
85
+ url = f"{BASE_URL}/{task_id}"
86
+
87
+ # 发起GET请求查询任务状态
88
+ response = requests.get(url, headers=HEADERS)
89
+
90
+ # 处理响应
91
+ if response.status_code == 200:
92
+ result = response.json()
93
+ task_status = result['data']['task_status']
94
+ print(f"任务状态: {task_status}")
95
+
96
+ # 如果任务成功,下载并保存生成的图片
97
+ if task_status == 'succeed':
98
+ images = result['data']['task_result']['images']
99
+ for image in images:
100
+ image_url = image['url']
101
+ image_index = i
102
+ save_image(image_url, image_index)
103
+
104
+ # 返回任务状态
105
+ return task_status
106
+ else:
107
+ print(f"查询任务失败: {response.status_code}, {response.text}")
108
+ return None
109
+
110
+ # ========== 步骤 3:下载并保存图片 ==========
111
+
112
+ def save_image(image_url, image_index):
113
+ """
114
+ 根据图片URL下载并保存到本地
115
+ """
116
+ response = requests.get(image_url)
117
+
118
+ if response.status_code == 200:
119
+ # 保存图片到本地
120
+ image_path = os.path.join(SAVE_DIRECTORY, f"kolor_{image_index}.png")
121
+ with open(image_path, "wb") as file:
122
+ file.write(response.content)
123
+ print(f"图片已保存到: {image_path}")
124
+ else:
125
+ print(f"下载图片失败: {response.status_code}, {response.text}")
126
+
127
+ # ========== 步骤 4:查询任务列表 ==========
128
+
129
+ def query_task_list(page_num=1, page_size=30):
130
+ """
131
+ 查询任务列表
132
+ """
133
+ # 查询参数
134
+ params = {
135
+ "pageNum": page_num,
136
+ "pageSize": page_size
137
+ }
138
+
139
+ # 发起GET请求查询任务列表
140
+ response = requests.get(BASE_URL, headers=HEADERS, params=params)
141
+
142
+ # 处理响应
143
+ if response.status_code == 200:
144
+ result = response.json()
145
+ print("任务列表:")
146
+ for task in result['data']:
147
+ print(f"任务ID: {task['task_id']}, 状态: {task['task_status']}, 创建时间: {task['created_at']}")
148
+ else:
149
+ print(f"查询任务列表失败: {response.status_code}, {response.text}")
150
+
151
+ # ========== 主流程 ==========
152
+
153
+ def kolor_vton(humen, cloth, i):
154
+ # 1. 创建虚���试穿任务
155
+ task_id = create_virtual_tryon_task(humen, cloth)
156
+
157
+ # 如果任务创建成功,继续执行查询步骤
158
+ if task_id:
159
+ # 2. 定期查询该任务状态,直到任务完成
160
+ while True:
161
+ status = query_virtual_tryon_task(task_id, i)
162
+ if status in ['succeed', 'failed']:
163
+ break
164
+ print("任务正在处理中,等待5秒后重试...")
165
+ time.sleep(5) # 等待5秒后重试
166
+
167
+ # 3. 查询任务列表
168
+ query_task_list(page_num=1, page_size=10)
169
+
170
+ # kolor_vton("uploads/user_image.jpg", "uploads/user_image.jpg", 1)