import requests
|
import urllib.parse
|
import datetime
|
import random
|
import os
|
from io import BytesIO
|
|
from auto_media_publisher.config.conf_base import (
|
TRANSLATE_API_TEMPLATE,
|
POLLINATIONS_API_TEMPLATE,
|
POLLINATIONS_DEFAULT_WIDTH ,
|
POLLINATIONS_DEFAULT_HEIGHT ,
|
POLLINATIONS_DEFAULT_MODEL,
|
)
|
|
|
class Text2ImageGenerator:
|
def __init__(self, width=POLLINATIONS_DEFAULT_WIDTH, height=POLLINATIONS_DEFAULT_HEIGHT, model=POLLINATIONS_DEFAULT_MODEL):
|
self.width = width
|
self.height = height
|
self.model = model
|
|
def translate_prompt(self, prompt: str) -> str:
|
"""
|
使用 MyMemory API 将中文提示词翻译为英文，失败则返回原文。
|
"""
|
|
encoded_prompt = urllib.parse.quote(prompt)
|
api_url = TRANSLATE_API_TEMPLATE.format(prompt=encoded_prompt)
|
try:
|
res = requests.get(api_url, timeout=30)
|
res.raise_for_status()
|
data = res.json()
|
if data.get("responseStatus") == 200 and "translatedText" in data.get("responseData", {}):
|
translated = data["responseData"]["translatedText"]
|
print(f"[✓] 翻译成功: {translated}")
|
return translated
|
except Exception as e:
|
print(f"[!] 翻译失败，使用原文: {e}")
|
return prompt
|
|
def build_url(self, prompt: str) -> str:
|
"""
|
构建 Pollinations 图片生成 URL。
|
"""
|
seed = random.randint(1, 1_000_000)
|
encoded_prompt = urllib.parse.quote(prompt)
|
return POLLINATIONS_API_TEMPLATE.format(
|
prompt=encoded_prompt,
|
width=self.width,
|
height=self.height,
|
seed=seed,
|
model=self.model
|
)
|
|
def _fetch_with_retry(self, url: str, max_retries: int = 3):
|
"""
|
最多尝试 max_retries 次请求 URL，成功则返回 response.content，否则返回 None。
|
"""
|
for attempt in range(1, max_retries + 1):
|
try:
|
response = requests.get(url, timeout=30)
|
response.raise_for_status()
|
return response.content
|
except Exception as e:
|
print(f"[!] 请求第 {attempt} 次失败：{e}")
|
return None
|
|
def generate_image(self, prompt: str, filename: str = None) -> str:
|
if not prompt.strip():
|
raise ValueError("提示词不能为空")
|
|
translated_prompt = self.translate_prompt(prompt)
|
url = self.build_url(translated_prompt)
|
|
content = self._fetch_with_retry(url)
|
if content is None:
|
return None
|
|
if not filename:
|
now = datetime.datetime.now()
|
filename = now.strftime("%Y%m%d%H%M%S") + ".png"
|
if not filename.lower().endswith(".png"):
|
filename += ".png"
|
|
with open(filename, "wb") as f:
|
f.write(content)
|
|
abs_path = os.path.abspath(filename)
|
print(f"[✓] 图片已保存：{abs_path}")
|
return abs_path
|
|
def generate_image_to_stream(self, prompt: str) -> BytesIO:
|
if not prompt.strip():
|
raise ValueError("提示词不能为空")
|
|
translated_prompt = self.translate_prompt(prompt)
|
url = self.build_url(translated_prompt)
|
|
content = self._fetch_with_retry(url)
|
if content is None:
|
return None
|
|
image_stream = BytesIO(content)
|
return image_stream
|
