1
+ import http .client
2
+ import json
3
+ import time
4
+ import base64
5
+
6
+ class KlingVirtualTryOn :
7
+ def __init__ (self , api_token , api_url ):
8
+ """初始化 Kling 虚拟试穿生成器
9
+
10
+ 参数:
11
+ api_token: API 密钥
12
+ api_url: API 节点地址
13
+ """
14
+ self .api_url = api_url
15
+ self .api_token = api_token
16
+ # 初始化 HTTP 连接
17
+ self .conn = http .client .HTTPSConnection (self .api_url )
18
+ self .endpoint = "/kling/v1/images/kolors-virtual-try-on"
19
+ # 设置请求头
20
+ self .headers = {
21
+ 'Authorization' : f'Bearer { self .api_token } ' ,
22
+ 'Content-Type' : 'application/json'
23
+ }
24
+
25
+ @staticmethod
26
+ def get_image_base64 (image_path ):
27
+ """将图片转换为 base64 编码形式
28
+
29
+ 参数:
30
+ image_path: 图片路径
31
+ 返回:
32
+ base64 编码后的图片字符串
33
+ """
34
+ with open (image_path , "rb" ) as image_file :
35
+ return base64 .b64encode (image_file .read ()).decode ("utf-8" )
36
+
37
+ def _kling_virtual_try_on (self , model_name , human_image , cloth_image , callback_url = "" ):
38
+ """使用 kling 生成虚拟试穿图像
39
+
40
+ 参数:
41
+ model_name: str, 模型名称,可选值:基础版本 `kolors-virtual-try-on-v1` 或 v1-5 版本 `kolors-virtual-try-on-v1-5` 支持服装组合
42
+ human_image: str, 人物图片(URL或base64编码)
43
+ cloth_image: str, 服饰图片(URL或base64编码)
44
+ callback_url: str, 回调地址,可以用于 webhook 等通知场景
45
+ 返回:
46
+ task_id: 生成任务的 id
47
+ """
48
+ # 构建请求体,请求的核心参数
49
+ payload = {
50
+ "model_name" : model_name ,
51
+ "human_image" : human_image ,
52
+ "cloth_image" : cloth_image ,
53
+ "callback_url" : callback_url
54
+ }
55
+
56
+ # 发送 POST 请求,提交虚拟试穿任务
57
+ self .conn .request ("POST" , self .endpoint , json .dumps (payload ), self .headers )
58
+ # 获取响应
59
+ res = self .conn .getresponse ()
60
+ # 读取响应内容并解析为 JSON
61
+ json_data = json .loads (res .read ().decode ("utf-8" ))
62
+
63
+ if 'code' in json_data and json_data ['code' ] == 0 :
64
+ # 成功则返回提交的任务 id
65
+ return json_data ['data' ]['task_id' ]
66
+ else :
67
+ # 失败则返回错误信息
68
+ raise Exception (f"API调用失败:{ json_data ['message' ]} " )
69
+
70
+ def _query_virtual_try_on_result (self , task_id ):
71
+ """使用查询接口获取虚拟试穿结果
72
+
73
+ 参数:
74
+ task_id: 生成任务的 id
75
+ 返回:
76
+ result_image: 虚拟试穿结果图像 url,任务未完成时返回 None
77
+ """
78
+ # 构建查询路径
79
+ query_path = f"{ self .endpoint } /{ task_id } "
80
+
81
+ # 发送 GET 请求,查询虚拟试穿任务状态
82
+ self .conn .request ("GET" , query_path , None , self .headers )
83
+ # 获取响应
84
+ res = self .conn .getresponse ()
85
+ # 读取响应内容并解析为 JSON
86
+ json_data = json .loads (res .read ().decode ("utf-8" ))
87
+
88
+ # 如果任务状态为成功,则返回结果图像 url
89
+ if json_data ['data' ]['task_status' ] == "succeed" :
90
+ result_image = json_data ['data' ]['task_result' ]['images' ][0 ]['url' ]
91
+ return result_image
92
+ else :
93
+ return None
94
+
95
+ def generate_try_on (self , model_name , human_image , cloth_image , callback_url = "" , timeout = 120 ):
96
+ """实现功能,根据人物图像和服饰图像生成虚拟试穿结果
97
+
98
+ 参数:
99
+ model_name: str, 模型名称,可选值:基础版本 `kolors-virtual-try-on-v1` 或 v1-5 版本 `kolors-virtual-try-on-v1-5` 支持服装组合
100
+ human_image: str, 人物图片路径或URL
101
+ cloth_image: str, 服饰图片路径或URL
102
+ callback_url: str, 回调地址,可以用于 webhook 等通知场景
103
+ timeout: int, 超时时间(秒)
104
+ 返回:
105
+ result_image: 虚拟试穿结果图像 url
106
+ """
107
+ # 处理人物图片输入
108
+ if human_image .startswith (('http://' , 'https://' , 'ftp://' )):
109
+ # 如果是URL,直接使用
110
+ human_data = human_image
111
+ else :
112
+ # 否则当作本地文件路径处理,转换为base64
113
+ try :
114
+ human_data = KlingVirtualTryOn .get_image_base64 (human_image )
115
+ except Exception as e :
116
+ raise ValueError (f"无法读取人物图像文件: { str (e )} " )
117
+
118
+ # 处理服饰图片输入
119
+ if cloth_image .startswith (('http://' , 'https://' , 'ftp://' )):
120
+ # 如果是URL,直接使用
121
+ cloth_data = cloth_image
122
+ else :
123
+ # 否则当作本地文件路径处理,转换为base64
124
+ try :
125
+ cloth_data = KlingVirtualTryOn .get_image_base64 (cloth_image )
126
+ except Exception as e :
127
+ raise ValueError (f"无法读取服饰图像文件: { str (e )} " )
128
+
129
+ # 调用虚拟试穿 API 提交任务,返回获取 task_id
130
+ task_id = self ._kling_virtual_try_on (model_name , human_data , cloth_data , callback_url )
131
+
132
+ start_time = time .time ()
133
+
134
+ # 轮询等待生成完成
135
+ while True :
136
+ # 根据 task_id 调用查询 API 查看任务是否完成
137
+ result_image = self ._query_virtual_try_on_result (task_id )
138
+ # 如果任务完成,则返回结果图像 url
139
+ if result_image is not None :
140
+ return result_image
141
+ # 如果轮询超时,则返回 None
142
+ if time .time () - start_time > timeout :
143
+ print (f"请求达到 { timeout } 秒超时" )
144
+ return None
145
+ # 轮询间隔 1 秒
146
+ time .sleep (1 )
147
+ print (f"等待虚拟试穿结果生成,{ int (time .time () - start_time )} 秒" , flush = True )
148
+
149
+
150
+ # 使用示例
151
+ if __name__ == "__main__" :
152
+ API_URL = "www.dmxapi.cn" # API 节点地址
153
+ DMX_API_TOKEN = "sk-XXXXXXXXXXXXXX" # API 密钥
154
+
155
+ # 创建虚拟试穿生成器实例
156
+ kling_virtual_try_on = KlingVirtualTryOn (api_token = DMX_API_TOKEN , api_url = API_URL )
157
+
158
+ # 生成虚拟试穿结果
159
+ result_url = kling_virtual_try_on .generate_try_on (
160
+ model_name = "kolors-virtual-try-on-v1-5" , # [必选] 模型名称 参数基础版本 `kolors-virtual-try-on-v1` 或 v1-5 版本 `kolors-virtual-try-on-v1-5` 支持服装组合
161
+ human_image = "https://assets.christiandior.com/is/image/diorprod/LOOK_F_25_1_LOOK_095_E04?$lookDefault_GH-GHC$&crop=568,0,1864,2000&bfc=on&qlt=85" , # [必选] 人物图片 参数可以是 图片的URL 或者 图片的本地路径 即可
162
+ cloth_image = "https://assets.christiandior.com/is/image/diorprod/511R59A1166X3389_E01?$default_GHC$&crop=501,147,998,1572&bfc=on&qlt=85" , # [必选] 服饰图片 参数可以是 图片的URL 或者 图片的本地路径 即可
163
+ # callback_url="", # 回调地址
164
+ # timeout=120 # 等待超时时间 绘图任务推荐设置 30 秒以上
165
+ )
166
+
167
+ print (result_url )
0 commit comments