'commit'

dsLightRag/QWenImage/QwenImageKit.py

@@ -17,21 +17,12 @@ class QwenImageGenerator:
         self.api_key = api_key or ALY_LLM_API_KEY
         self.model = "qwen-image"
         
-    def generate_image(self, prompt, n=1, size='1328*1328'):
+    def generate_image(self, prompt, size='1328*1328'):  # 移除n参数
         """生成图片
         
         Args:
             prompt: 图片描述提示词
-            n: 生成图片数量
             size: 图片尺寸
-            
-        Returns:
-            dict: 包含生成结果的字典，格式如下：
-                {
-                    'success': bool,  # 操作是否成功
-                    'images': list,   # 图片URL列表
-                    'error_msg': str  # 错误信息，如果成功则为None
-                }
         """
         try:
             # 调用文心一言图片生成API
@@ -39,7 +30,6 @@ class QwenImageGenerator:
                 api_key=self.api_key,
                 model=self.model,
                 prompt=prompt,
-                n=n,
                 size=size
             )
             
@@ -47,6 +37,13 @@ class QwenImageGenerator:
             if rsp.status_code == HTTPStatus.OK:
                 # 提取图片URL
                 image_urls = [result.url for result in rsp.output.results]
+                # 修改错误消息，移除n参数引用
+                if not image_urls:
+                    return {
+                        'success': False,
+                        'images': [],
+                        'error_msg': 'API返回空结果，图片生成失败'
+                    }
                 return {
                     'success': True,
                     'images': image_urls,
@@ -97,7 +94,7 @@ class QwenImageGenerator:
                 'error_msg': f'保存图片失败: {str(e)}'
             }
     
-    def generate_and_save_images(self, prompt, n=1, size='1328*1328', save_dir='./'):
+    def generate_and_save_images(self, prompt, size='1328*1328', save_dir='./'):
         """生成图片并保存到本地
         
         Args:
@@ -110,7 +107,7 @@ class QwenImageGenerator:
             dict: 包含生成和保存结果的字典
         """
         # 生成图片
-        generate_result = self.generate_image(prompt, n, size)
+        generate_result = self.generate_image(prompt,  size)
         
         if not generate_result['success']:
             return generate_result