fix: kling text to image

Author: timerring

docs/zh/models/kling/kling-text-to-image.md

@@ -24,40 +24,41 @@ updatedAt: 2025-05-08
 > 深色背景为可以修改的参数，非必选参数已经注释，可以按照自己的需求启用。
 
 
-::: code-group
-
-
-<<< @/zh/snippets/kling-text-to-image.py{121-122,129-135}[标准类版本]
-
-
-
-<<< @/zh/snippets/kling-text-to-image-function.py{6-7,20-26}[纯函数版本]
-
-
-
-:::
+<<< @/zh/snippets/kling-text-to-image.py{120-121,128-135}
 
 
 ## 返回结果
 
-返回结果为图片的 url，每个 url 有效期一般为 3 天以内，尽快下载或者转存。
+返回结果为图片的 url 列表，这里使用参数 `n=2` 生成两张图片，每个 url kling 官方保证有效期为 30 天，尽快下载或者转存。
 
 ```
-https://cdn.klingai.com/bs2/upload-kling-api/6567899185/image/CjikMGgHQaYAAAAAA1CAQw-0_raw_image_0.png
+['https://cdn.klingai.com/bs2/upload-kling-api/8089468206/image/Cl6kH2gHPegAAAAABUwweg-0_raw_image_0.png', 'https://cdn.klingai.com/bs2/upload-kling-api/8089468206/image/Cl6kH2gHPegAAAAABUwweg-1_raw_image_0.png']
 ```
 
-![](https://cdn.jsdelivr.net/gh/timerring/scratchpad2023/2024/2025-05-04-23-33-40.png)
+![](https://cdn.jsdelivr.net/gh/timerring/scratchpad2023/2024/2025-05-11-20-50-49.png)
+
+![](https://cdn.jsdelivr.net/gh/timerring/scratchpad2023/2024/2025-05-11-20-51-02.png)
 
 ## 函数流程图
 
 ```mermaid
-flowchart TD
-J[调用kling_generate_image获取task_id] --> K[开始轮询任务状态]
-K --> L[调用query_kling_image_url]
-L --> M{获取到图像URL?}
-M -->|否| N{是否超时?}
-N -->|否| O[等待1秒]
-O --> L
-N -->|是| P[返回超时错误]
-M -->|是| Q[返回图像URL]
+graph TD
+    A[客户端初始化 KlingTextToImage] -->|提供 API 凭证| B[调用 generate_image 方法]
+    B -->|传入提示词和参数| C[_kling_generate_image 方法]
+    C -->|构建请求体| D[发送 POST 请求到 API]
+    D -->|返回 task_id| E[开始轮询任务状态]
+    
+    E -->|调用| F[_query_kling_image_url 方法]
+    F -->|发送 GET 请求| G[检查任务状态]
+    
+    G -->|任务状态?| H{任务完成?}
+    H -->|是| I[提取并返回图像 URL]
+    H -->|否| J{是否超时?}
+    
+    J -->|是| K[返回 None]
+    J -->|否| L[等待 1 秒]
+    L --> F
+    
+    I -->|返回图像 URL 列表| M[结束流程]
+    K -->|返回超时信息| M
 ```

docs/zh/snippets/kling-text-to-image-function.py

@@ -20,7 +20,7 @@ def __init__(self, api_token, api_url):
             'Content-Type': 'application/json'
         }
 
-    def kling_generate_image(self, model_name, prompt, negative_prompt, output_format, n, aspect_ratio, callback_url):
+    def _kling_generate_image(self, model_name, prompt, negative_prompt, output_format, n, aspect_ratio, callback_url):
         """使用 kling 生成图像
         
         参数: 
@@ -55,7 +55,7 @@ def kling_generate_image(self, model_name, prompt, negative_prompt, output_forma
         # 从 json 文件中返回任务 ID
         return json_data['data']['task_id']
 
-    def query_kling_image_url(self, task_id):
+    def _query_kling_image_url(self, task_id):
         """使用查询接口获取生成图像 url
         
         输入参数:
@@ -77,11 +77,12 @@ def query_kling_image_url(self, task_id):
         json_data = json.loads(res.read().decode("utf-8"))
         # 如果任务状态为成功，则返回图像 url
         if json_data['data']['task_status'] == "succeed":
-            return json_data['data']['task_result']['images'][0]['url']
+            image_urls = [image['url'] for image in json_data['data']['task_result']['images']]
+            return image_urls
         else: 
             return None
 
-    def generate_image(self, model_name, prompt, negative_prompt="", output_format="png", n=1, aspect_ratio="16:9", callback_url=""):
+    def generate_image(self, model_name, prompt, negative_prompt="", output_format="png", n=1, aspect_ratio="16:9", callback_url="", timeout=60):
         """实现功能，直接根据预设的参数返回生成图像的 url
         
         参数:
@@ -96,14 +97,12 @@ def generate_image(self, model_name, prompt, negative_prompt="", output_format="
             image_url: 图像 url
         """
         # 调用生成图像 api 提交图像生成任务，返回获取 task_id。
-        task_id = self.kling_generate_image(model_name, prompt, negative_prompt, output_format, n, aspect_ratio, callback_url) 
+        task_id = self._kling_generate_image(model_name, prompt, negative_prompt, output_format, n, aspect_ratio, callback_url) 
         start_time = time.time()
-        # 队列等待超时时间
-        timeout = 60
         # 轮询等待生成完成
         while True:
             # 根据 task_id 调用查询图像api 查看图像生成任务是否完成。
-            image_url = self.query_kling_image_url(task_id) 
+            image_url = self._query_kling_image_url(task_id) 
             # 如果图像生成任务完成，则返回图像 url
             if image_url is not None:
                 return image_url
@@ -119,7 +118,7 @@ def generate_image(self, model_name, prompt, negative_prompt="", output_format="
 # 使用示例
 if __name__ == "__main__":
     API_URL="www.dmxapi.cn" # API 节点地址
-    DMX_API_TOKEN = "sk-XXXXXXXXXXX"  # API 密钥
+    DMX_API_TOKEN = "sk-XXXXXXXXXXXXXX"  # API 密钥
 
     # 创建图像生成器实例
     kling_text_to_image = KlingTextToImage(api_token=DMX_API_TOKEN, api_url=API_URL)
@@ -132,7 +131,8 @@ def generate_image(self, model_name, prompt, negative_prompt="", output_format="
         # output_format="png", # 输出格式：png 或 jpg
         # n=1, # int, 生成数量 [1, 9]
         # aspect_ratio="16:9", # 输出比例：16:9, 9:16, 1:1, 4:3, 3:4, 3:2, 2:3
-        # callback_url="" # 回调地址，可以用于 webhook 等通知场景
+        # callback_url="", # 回调地址，可以用于 webhook 等通知场景
+        # timeout=120 # 队列等待超时时间
     )
 
     print(image_url)