'commit'

dsLightRag/QWenImage/QWenImageEditKit.py

@@ -0,0 +1,167 @@
+from http import HTTPStatus
+from urllib.parse import urlparse, unquote
+from pathlib import PurePosixPath
+import requests
+from dashscope import MultiModalConversation
+from Config.Config import ALY_LLM_API_KEY
+
+class QwenImageEditor:
+    """Qwen Image Editing Toolkit"""
+    def __init__(self, api_key=None):
+        """初始化图片编辑器
+        
+        Args:
+            api_key: 阿里云API密钥，如果为None则使用配置文件中的密钥
+        """
+        self.api_key = api_key or ALY_LLM_API_KEY
+        self.model = "qwen-image-edit"
+        
+    def edit_image(self, image_url, prompt, negative_prompt="", stream=False, watermark=True):
+        """编辑图片
+        
+        Args:
+            image_url: 原始图片URL
+            prompt: 编辑指令描述
+            negative_prompt: 负面提示词
+            stream: 是否流式返回
+            watermark: 是否添加水印
+        
+        Returns:
+            dict: 包含编辑结果的字典
+        """
+        try:
+            # 构建对话消息
+            messages = [{
+                "role": "user",
+                "content": [
+                    {"image": image_url},
+                    {"text": prompt}
+                ]
+            }]
+            
+            # 调用图片编辑API
+            response = MultiModalConversation.call(
+                api_key=self.api_key,
+                model=self.model,
+                messages=messages,
+                result_format='message',
+                stream=stream,
+                watermark=watermark,
+                negative_prompt=negative_prompt
+            )
+            
+            # 处理响应结果
+            if response.status_code == HTTPStatus.OK:
+                # 提取编辑后的图片URL
+                edited_image_url = response.output.choices[0].message.content[0].get("image")
+                
+                if not edited_image_url:
+                    return {
+                        'success': False,
+                        'image_url': None,
+                        'error_msg': 'API返回空结果，图片编辑失败'
+                    }
+                
+                return {
+                    'success': True,
+                    'image_url': edited_image_url,
+                    'error_msg': None
+                }
+            else:
+                error_msg = f'API调用失败: status_code={response.status_code}, code={response.code}, message={response.message}'
+                return {
+                    'success': False,
+                    'image_url': None,
+                    'error_msg': error_msg
+                }
+        except Exception as e:
+            return {
+                'success': False,
+                'image_url': None,
+                'error_msg': f'发生异常: {str(e)}'
+            }
+    
+    def save_edited_image(self, image_url, save_dir='./'):
+        """保存编辑后的图片到本地
+        
+        Args:
+            image_url: 编辑后的图片URL
+            save_dir: 保存目录
+        
+        Returns:
+            dict: 包含保存结果的字典
+        """
+        try:
+            # 从URL解析文件名
+            file_name = PurePosixPath(unquote(urlparse(image_url).path)).parts[-1]
+            save_path = f'{save_dir}{file_name}'
+            
+            # 下载并保存图片
+            with open(save_path, 'wb+') as f:
+                f.write(requests.get(image_url).content)
+            
+            return {
+                'success': True,
+                'file_path': save_path,
+                'error_msg': None
+            }
+        except Exception as e:
+            return {
+                'success': False,
+                'file_path': None,
+                'error_msg': f'保存图片失败: {str(e)}'
+            }
+    
+    def edit_and_save_image(self, image_url, prompt, save_dir='./', negative_prompt=""):
+        """编辑图片并保存到本地
+        
+        Args:
+            image_url: 原始图片URL
+            prompt: 编辑指令描述
+            save_dir: 保存目录
+            negative_prompt: 负面提示词
+        
+        Returns:
+            dict: 包含编辑和保存结果的字典
+        """
+        # 编辑图片
+        edit_result = self.edit_image(image_url, prompt, negative_prompt)
+        
+        if not edit_result['success']:
+            return edit_result
+        
+        # 保存编辑后的图片
+        save_result = self.save_edited_image(edit_result['image_url'], save_dir)
+        
+        if save_result['success']:
+            return {
+                'success': True,
+                'image_url': edit_result['image_url'],
+                'file_path': save_result['file_path'],
+                'error_msg': None
+            }
+        else:
+            return {
+                'success': False,
+                'image_url': edit_result['image_url'],
+                'file_path': None,
+                'error_msg': save_result['error_msg']
+            }
+
+# 示例使用
+if __name__ == "__main__":
+    # 创建编辑器实例
+    editor = QwenImageEditor()
+    
+    # 示例参数
+    image_url = "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"
+    edit_prompt = "将图中的人物改为站立姿势，弯腰握住狗的前爪"
+    
+    print('----编辑图片，请等待任务执行----')
+    result = editor.edit_and_save_image(image_url, edit_prompt)
+    
+    if result['success']:
+        print(f'编辑成功，图片URL: {result["image_url"]}')
+        print(f'保存成功，文件路径: {result["file_path"]}')
+    else:
+        print(f'编辑失败: {result["error_msg"]}')