220 lines
9.9 KiB
Python
220 lines
9.9 KiB
Python
import os
|
||
import time
|
||
import json
|
||
import requests
|
||
from dotenv import load_dotenv
|
||
import dashscope
|
||
from dashscope import ImageSynthesis, Generation
|
||
|
||
# Load environment variables
|
||
load_dotenv()
|
||
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
|
||
|
||
class ImageGenerator:
|
||
def __init__(self, provider="dashscope", model=None):
|
||
self.provider = provider
|
||
self.model = model
|
||
self.api_key = None
|
||
|
||
if provider == "doubao":
|
||
self.api_key = os.getenv("volcengine_API_KEY")
|
||
if not self.model:
|
||
self.model = "doubao-seedream-5-0-260128" # Default model from user input
|
||
elif provider == "dashscope":
|
||
self.api_key = os.getenv("DASHSCOPE_API_KEY")
|
||
if not self.model:
|
||
self.model = "wanx2.0-t2i-turbo"
|
||
|
||
def optimize_prompt(self, asr_text, progress_callback=None):
|
||
"""Use LLM to optimize the prompt"""
|
||
print(f"Optimizing prompt for: {asr_text}")
|
||
|
||
if progress_callback:
|
||
progress_callback(0, "正在准备优化提示词...")
|
||
|
||
system_prompt = """你是一个AI图像提示词优化专家。你的任务是将用户的语音识别结果转化为适合生成"黑白线稿"的提示词。
|
||
关键要求：
|
||
1. 风格必须是：简单的黑白线稿、简笔画、图标风格 (Line art, Sketch, Icon style)。
|
||
2. 画面必须清晰、线条粗壮，适合低分辨率热敏打印机打印。
|
||
3. 绝对不要有复杂的阴影、渐变、黑白线条描述。
|
||
4. 背景必须是纯白 (White background)。
|
||
5. 提示词内容请使用英文描述，因为绘图模型对英文理解更好，但在描述中强调 "black and white line art", "simple lines", "vector style"。
|
||
6. 尺寸比例遵循宽48mm:高30mm (约 1.6:1)。
|
||
7. 直接输出优化后的提示词，不要包含任何解释。
|
||
如果用户要求输入文字，则用```把文字包裹起来，文字是中文
|
||
"black and white line art， Chinese:```中国人```"
|
||
"""
|
||
|
||
try:
|
||
if progress_callback:
|
||
progress_callback(10, "正在调用AI优化提示词...")
|
||
|
||
# Currently using Qwen-Turbo for all providers for prompt optimization
|
||
# You can also decouple this if needed
|
||
response = Generation.call(
|
||
model='qwen-turbo',
|
||
prompt=f'{system_prompt}\n\n用户语音识别结果：{asr_text}\n\n优化后的提示词：',
|
||
max_tokens=200,
|
||
temperature=0.8
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
if hasattr(response, 'output') and response.output and \
|
||
hasattr(response.output, 'choices') and response.output.choices and \
|
||
len(response.output.choices) > 0:
|
||
|
||
optimized = response.output.choices[0].message.content.strip()
|
||
print(f"Optimized prompt: {optimized}")
|
||
|
||
if progress_callback:
|
||
progress_callback(30, f"提示词优化完成: {optimized[:50]}...")
|
||
|
||
return optimized
|
||
elif hasattr(response, 'output') and response.output and hasattr(response.output, 'text'):
|
||
optimized = response.output.text.strip()
|
||
print(f"Optimized prompt (direct text): {optimized}")
|
||
if progress_callback:
|
||
progress_callback(30, f"提示词优化完成: {optimized[:50]}...")
|
||
return optimized
|
||
else:
|
||
print(f"Prompt optimization response format error: {response}")
|
||
if progress_callback:
|
||
progress_callback(0, "提示词优化响应格式错误")
|
||
return asr_text
|
||
else:
|
||
print(f"Prompt optimization failed: {response.code} - {response.message}")
|
||
if progress_callback:
|
||
progress_callback(0, f"提示词优化失败: {response.message}")
|
||
return asr_text
|
||
|
||
except Exception as e:
|
||
print(f"Error optimizing prompt: {e}")
|
||
if progress_callback:
|
||
progress_callback(0, f"提示词优化出错: {str(e)}")
|
||
return asr_text
|
||
|
||
def generate_image(self, prompt, progress_callback=None):
|
||
"""Generate image based on provider"""
|
||
if self.provider == "dashscope":
|
||
return self._generate_dashscope(prompt, progress_callback)
|
||
elif self.provider == "doubao":
|
||
return self._generate_doubao(prompt, progress_callback)
|
||
else:
|
||
raise ValueError(f"Unknown provider: {self.provider}")
|
||
|
||
def _generate_dashscope(self, prompt, progress_callback=None):
|
||
print(f"Generating image with DashScope for prompt: {prompt}")
|
||
|
||
if progress_callback:
|
||
progress_callback(35, "正在请求DashScope生成图片...")
|
||
|
||
try:
|
||
response = ImageSynthesis.call(
|
||
model=self.model,
|
||
prompt=prompt,
|
||
size='1280*720'
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
if not response.output:
|
||
print("Error: response.output is None")
|
||
return None
|
||
|
||
task_status = response.output.get('task_status')
|
||
|
||
if task_status == 'PENDING' or task_status == 'RUNNING':
|
||
print("Waiting for image generation to complete...")
|
||
if progress_callback:
|
||
progress_callback(45, "AI正在生成图片中...")
|
||
|
||
task_id = response.output.get('task_id')
|
||
max_wait = 120
|
||
waited = 0
|
||
while waited < max_wait:
|
||
time.sleep(2)
|
||
waited += 2
|
||
task_result = ImageSynthesis.fetch(task_id)
|
||
if task_result.output.task_status == 'SUCCEEDED':
|
||
response.output = task_result.output
|
||
break
|
||
elif task_result.output.task_status == 'FAILED':
|
||
error_msg = task_result.output.message if hasattr(task_result.output, 'message') else 'Unknown error'
|
||
print(f"Image generation failed: {error_msg}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成失败: {error_msg}")
|
||
return None
|
||
|
||
if response.output.get('task_status') == 'SUCCEEDED':
|
||
image_url = response.output['results'][0]['url']
|
||
print(f"Image generated, url: {image_url}")
|
||
return image_url
|
||
else:
|
||
error_msg = f"{response.code} - {response.message}"
|
||
print(f"Image generation failed: {error_msg}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成失败: {error_msg}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"Error generating image: {e}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成出错: {str(e)}")
|
||
return None
|
||
|
||
def _generate_doubao(self, prompt, progress_callback=None):
|
||
print(f"Generating image with Doubao for prompt: {prompt}")
|
||
|
||
if progress_callback:
|
||
progress_callback(35, "正在请求豆包生成图片...")
|
||
|
||
url = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {self.api_key}"
|
||
}
|
||
data = {
|
||
"model": self.model,
|
||
"prompt": prompt,
|
||
"sequential_image_generation": "disabled",
|
||
"response_format": "url",
|
||
"size": "2K", # Doubao supports different sizes, user example used 2K. But we might want something smaller if possible to save bandwidth/time?
|
||
# User's curl says "2K". I will stick to it or maybe check docs.
|
||
# Actually for thermal printer, we don't need 2K. But let's follow user example.
|
||
"stream": False,
|
||
"watermark": True
|
||
}
|
||
|
||
try:
|
||
response = requests.post(url, headers=headers, json=data, timeout=60)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
# Check format of result
|
||
# Typically OpenAI compatible or similar
|
||
# User example didn't show response format, but usually it's "data": [{"url": "..."}]
|
||
|
||
if "data" in result and len(result["data"]) > 0:
|
||
image_url = result["data"][0]["url"]
|
||
print(f"Image generated, url: {image_url}")
|
||
return image_url
|
||
elif "error" in result:
|
||
error_msg = result["error"].get("message", "Unknown error")
|
||
print(f"Doubao API error: {error_msg}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成失败: {error_msg}")
|
||
return None
|
||
else:
|
||
print(f"Unexpected response format: {result}")
|
||
return None
|
||
else:
|
||
print(f"Doubao API failed with status {response.status_code}: {response.text}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成失败: {response.status_code}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"Error calling Doubao API: {e}")
|
||
if progress_callback:
|
||
progress_callback(35, f"图片生成出错: {str(e)}")
|
||
return None
|