194 lines
9.0 KiB
Python
194 lines
9.0 KiB
Python
import logging
|
||
import uuid
|
||
import asyncio
|
||
from datetime import datetime
|
||
from fastapi import APIRouter, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from pydantic import BaseModel
|
||
from sqlalchemy.sql import text
|
||
|
||
from Util.BananaClient import BananaClient
|
||
from Util.LlmUtil import get_llm_response
|
||
from Model.HaiBaoModel import HaiBaoModel
|
||
|
||
router = APIRouter(prefix="/haibao")
|
||
logger = logging.getLogger("HaiBaoController")
|
||
|
||
class GenerateRequest(BaseModel):
|
||
prompt: str
|
||
width: int = 1024
|
||
height: int = 1024
|
||
|
||
@router.on_event("startup")
|
||
async def startup_event():
|
||
"""初始化时逻辑"""
|
||
logger.info("海报控制器启动成功")
|
||
|
||
class RefineRequest(BaseModel):
|
||
prompt: str
|
||
|
||
@router.post("/refine")
|
||
async def refine_prompt(req: RefineRequest):
|
||
"""润色提示词"""
|
||
async def generate():
|
||
try:
|
||
extra_requirements = ""
|
||
if "春节" in req.prompt:
|
||
extra_requirements = "7. 画面氛围要极其温馨喜庆，体现春节团圆主题。可以尝试在画面中自然融入“过年了”、“回家看看”等字样（如果模型支持文字生成），或者通过红灯笼、春联、全家福等元素强烈暗示这一主题。"
|
||
|
||
refine_system_prompt = "你是一个资深的AI绘画提示词专家。你的任务是将用户简短的描述扩充为一段详细、高质量的画面描述提示词，用于生成宣传海报。"
|
||
refine_user_prompt = f"""
|
||
请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
|
||
|
||
主题：{req.prompt}
|
||
|
||
要求：
|
||
1. 描述画面主体、背景、光影、色彩、构图。
|
||
2. 风格要求：现代感、科技感、精美、3D渲染风格或高品质插画风格。
|
||
3. 融入新能源、绿色环保、充电桩等元素。
|
||
4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
|
||
5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
|
||
6. 字数在100-300字之间。
|
||
{extra_requirements}
|
||
"""
|
||
|
||
async for chunk in get_llm_response(query_text=refine_user_prompt, system_prompt=refine_system_prompt, stream=True):
|
||
yield chunk
|
||
except Exception as e:
|
||
logger.error(f"提示词润色失败: {e}")
|
||
yield f"Error: {str(e)}"
|
||
|
||
return StreamingResponse(generate(), media_type="text/plain")
|
||
|
||
@router.post("/generate")
|
||
async def generate_poster(req: GenerateRequest):
|
||
"""生成海报及文案"""
|
||
try:
|
||
# 并行执行生图和生文
|
||
client = BananaClient()
|
||
|
||
# 1. 构造生图任务 (包含智能润色判断)
|
||
async def generate_image_task():
|
||
final_prompt = req.prompt
|
||
|
||
# 智能判断：如果提示词太短（少于50字），则认为用户未进行润色，自动执行润色
|
||
# 如果用户使用了"一键扩写"功能，提示词通常会很长，这里就会跳过自动润色，尊重用户的修改
|
||
if len(final_prompt) < 50:
|
||
logger.info(f"提示词较短({len(final_prompt)}字)，执行自动润色...")
|
||
|
||
extra_requirements = ""
|
||
if "春节" in req.prompt:
|
||
extra_requirements = "7. 画面氛围要极其温馨喜庆，体现春节团圆主题。可以尝试在画面中自然融入“过年了”、“回家看看”等字样（如果模型支持文字生成），或者通过红灯笼、春联、全家福等元素强烈暗示这一主题。"
|
||
|
||
refine_system_prompt = "你是一个资深的AI绘画提示词专家。你的任务是将用户简短的描述扩充为一段详细、高质量的画面描述提示词，用于生成宣传海报。"
|
||
refine_user_prompt = f"""
|
||
请根据以下主题，为充电企业“驿来特”设计一张宣传海报的画面描述。
|
||
|
||
主题：{req.prompt}
|
||
|
||
要求：
|
||
1. 描述画面主体、背景、光影、色彩、构图。
|
||
2. 风格要求：现代感、科技感、精美、3D渲染风格或高品质插画风格。
|
||
3. 融入新能源、绿色环保、充电桩等元素。
|
||
4. 关于品牌元素：画面中可自然融入品牌Logo的视觉元素（如配色、形状），能用多少就用多少，有元素体现即可，不必生搬硬套，保持画面自然和谐。
|
||
5. 直接输出提示词内容，不要包含“好的”、“以下是”等无关废话。
|
||
6. 字数在100-300字之间。
|
||
{extra_requirements}
|
||
"""
|
||
|
||
refined_prompt = ""
|
||
try:
|
||
async for chunk in get_llm_response(query_text=refine_user_prompt, system_prompt=refine_system_prompt, stream=False):
|
||
refined_prompt += chunk
|
||
if refined_prompt and refined_prompt.strip():
|
||
final_prompt = refined_prompt
|
||
except Exception as e:
|
||
logger.error(f"自动润色失败，使用原始提示词: {e}")
|
||
|
||
logger.info(f"Final generation prompt: {final_prompt}")
|
||
|
||
# 1.2 调用生图
|
||
resp = await client.generate_image(prompt=final_prompt, size=f"{req.width}x{req.height}")
|
||
# 定义Logo路径
|
||
LOGO_PATH = r"d:\dsWork\aiData\static\Images\login_logo.png"
|
||
obs_urls = await client.download_and_upload_to_obs(resp, overlay_logo_path=LOGO_PATH)
|
||
if not obs_urls:
|
||
raise Exception("未获取到有效的图片URL")
|
||
return obs_urls[0]
|
||
|
||
# 2. 构造生文任务
|
||
async def generate_text_task():
|
||
extra_text_req = ""
|
||
if "春节" in req.prompt:
|
||
extra_text_req = "6. 文案中必须包含“过年了，回家看看，记得提前充电噢”或类似的温馨提示。"
|
||
|
||
scheme_prompt = f"""
|
||
你是一个专业的社群运营专家。请为充电企业“驿来特”撰写一段发在微信群里的宣传文案。
|
||
|
||
主题：{req.prompt}
|
||
|
||
要求：
|
||
1. 语气亲切、有吸引力，适合微信社群传播。
|
||
2. 突出“驿来特”品牌，强调新能源、优惠、便利等特点（根据主题自由发挥）。
|
||
3. 包含适当的emoji表情，增加趣味性。
|
||
4. 字数控制在150字以内。
|
||
5. 格式清晰，分段合理。
|
||
{extra_text_req}
|
||
"""
|
||
# get_llm_response 是一个异步生成器 (stream=True by default) 或者直接返回 (stream=False)
|
||
# 这里我们强制 stream=False 获取完整文本
|
||
text_response = ""
|
||
# LlmUtil.get_llm_response 默认为 stream=True，我们需要修改调用方式或适配
|
||
# 查看 LlmUtil 源码，如果 stream=False，它 yield 内容。
|
||
# 所以我们需要迭代它
|
||
async for chunk in get_llm_response(query_text=scheme_prompt, stream=False):
|
||
text_response += chunk
|
||
return text_response
|
||
|
||
# 3. 并行执行
|
||
image_url, scheme_content = await asyncio.gather(generate_image_task(), generate_text_task())
|
||
|
||
# 4. 保存到数据库
|
||
model = HaiBaoModel()
|
||
record_id = str(uuid.uuid4())
|
||
created_at = datetime.now()
|
||
|
||
await model.insert_record(record_id, req.prompt, image_url, scheme_content, created_at)
|
||
|
||
return {
|
||
"id": record_id,
|
||
"image_url": image_url,
|
||
"prompt": req.prompt,
|
||
"scheme_content": scheme_content,
|
||
"created_at": created_at.strftime("%Y-%m-%d %H:%M:%S")
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"生成海报/文案失败: {e}")
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
@router.get("/history")
|
||
async def get_history():
|
||
"""获取海报生成历史"""
|
||
try:
|
||
model = HaiBaoModel()
|
||
result = await model.get_history(50)
|
||
|
||
formatted_result = []
|
||
for item in result:
|
||
item_dict = dict(item) if not isinstance(item, dict) else item
|
||
|
||
if isinstance(item_dict.get('created_at'), datetime):
|
||
item_dict['created_at'] = item_dict['created_at'].strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
# 确保 scheme_content 存在
|
||
if 'scheme_content' not in item_dict:
|
||
item_dict['scheme_content'] = ""
|
||
|
||
formatted_result.append(item_dict)
|
||
|
||
return formatted_result
|
||
except Exception as e:
|
||
logger.error(f"获取历史失败: {e}")
|
||
return []
|