dsProject/dsLightRag/QWenImage/QwenImageKit.py

from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
import requests
# pip install -U dashscope
from dashscope import ImageSynthesis
from Config.Config import ALY_LLM_API_KEY

class QwenImageGenerator:

    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"
        
    def generate_image(self, prompt, size='1328*1328'):  # 移除n参数
        """生成图片
        
        Args:
            prompt: 图片描述提示词
            size: 图片尺寸
        """
        try:
            # 调用文心一言图片生成API
            rsp = ImageSynthesis.call(
                api_key=self.api_key,
                model=self.model,
                prompt=prompt,
                size=size
            )
            
            # 处理响应结果
            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,
                    'error_msg': None
                }
            else:
                error_msg = f'API调用失败: status_code={rsp.status_code}, code={rsp.code}, message={rsp.message}'
                return {
                    'success': False,
                    'images': [],
                    'error_msg': error_msg
                }
        except Exception as e:
            return {
                'success': False,
                'images': [],
                'error_msg': f'发生异常: {str(e)}'
            }
    
    def save_image_from_url(self, image_url, save_dir='./'):
        """从URL保存图片到本地
        
        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 generate_and_save_images(self, prompt, size='1328*1328', save_dir='./'):
        """生成图片并保存到本地
        
        Args:
            prompt: 图片描述提示词
            n: 生成图片数量
            size: 图片尺寸
            save_dir: 保存目录
            
        Returns:
            dict: 包含生成和保存结果的字典
        """
        # 生成图片
        generate_result = self.generate_image(prompt,  size)
        
        if not generate_result['success']:
            return generate_result
        
        # 保存所有生成的图片
        saved_files = []
        for image_url in generate_result['images']:
            save_result = self.save_image_from_url(image_url, save_dir)
            if save_result['success']:
                saved_files.append(save_result['file_path'])
            else:
                return {
                    'success': False,
                    'images': generate_result['images'],
                    'saved_files': saved_files,
                    'error_msg': save_result['error_msg']
                }
        
        return {
            'success': True,
            'images': generate_result['images'],
            'saved_files': saved_files,
            'error_msg': None
        }

# 示例使用
if __name__ == "__main__":
    # 创建生成器实例
    generator = QwenImageGenerator()
    
    # 示例提示词
    prompt = """一副典雅庄重的对联悬挂于厅堂之中，房间是个安静古典的中式布置，桌子上放着一些青花瓷，对联上左书"云端书声琅琅跨山海"，右书"智脉笔影翩翩启星辰"， 横批"东师理想"，字体飘逸，中间挂在一着一副中国风的画作，内容是岳阳楼。"""
    
    print('----生成图片，请等待任务执行----')
    result = generator.generate_and_save_images(prompt)
    
    if result['success']:
        print(f'生成成功，图片URL: {result["images"]}')
        print(f'保存成功，文件路径: {result["saved_files"]}')
    else:
        print(f'生成失败: {result["error_msg"]}')