dsProject/dsLightRag/QWenImage/QWenImageEditKit.py

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"]}')