theliu
216 lines
7.0 KiB
Python
216 lines
7.0 KiB
Python
"""
|
||
image_gen.py - 统一文生图接口
|
||
支持两个模型：
|
||
- Kolors（便宜快速）→ SiliconFlow API（同步）
|
||
- Qwen-Image（高质量）→ ModelScope API（异步轮询）
|
||
"""
|
||
|
||
import requests
|
||
import os
|
||
import time
|
||
from datetime import datetime
|
||
from config import (
|
||
SILICONFLOW_API_KEY,
|
||
SILICONFLOW_API_BASE,
|
||
MODELSCOPE_API_KEY,
|
||
MODELSCOPE_API_BASE,
|
||
MODELSCOPE_POLL_INTERVAL,
|
||
MODELSCOPE_MAX_WAIT,
|
||
IMAGE_MODELS,
|
||
NEGATIVE_PROMPT,
|
||
)
|
||
|
||
|
||
def _generate_siliconflow(prompt, model_id, size, guidance, neg, save_dir, filename):
|
||
"""SiliconFlow 同步 API（Kolors）"""
|
||
payload = {
|
||
"model": model_id,
|
||
"prompt": prompt,
|
||
"image_size": size,
|
||
"n": 1,
|
||
"num_inference_steps": 20,
|
||
"guidance_scale": guidance,
|
||
"negative_prompt": neg,
|
||
}
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {SILICONFLOW_API_KEY}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
print(f" [SiliconFlow] 提交: {prompt[:60]}{'...' if len(prompt) > 60 else ''}")
|
||
|
||
for attempt in range(6): # 最多重试 5 次
|
||
resp = requests.post(SILICONFLOW_API_BASE, headers=headers, json=payload, timeout=120)
|
||
print(f" HTTP {resp.status_code}: {resp.text[:300]}")
|
||
|
||
if resp.status_code == 429:
|
||
wait = 15 * (attempt + 1) # 15s, 30s, 45s, 60s, 75s
|
||
print(f" [!] 限频，等待 {wait}s 后重试 ({attempt+1}/5)...")
|
||
time.sleep(wait)
|
||
continue
|
||
|
||
if resp.status_code != 200:
|
||
raise Exception(f"SiliconFlow 生成失败 ({resp.status_code}): {resp.text[:300]}")
|
||
break
|
||
else:
|
||
raise Exception("SiliconFlow 持续限频，已重试 5 次，请稍后再试或切换模型")
|
||
|
||
result = resp.json()
|
||
images = result.get("images", [])
|
||
if not images:
|
||
raise Exception(f"SiliconFlow 返回无图片: {result}")
|
||
|
||
img_url = images[0].get("url")
|
||
if not img_url:
|
||
raise Exception(f"返回图片 URL 为空: {result}")
|
||
|
||
img_data = requests.get(img_url, timeout=60).content
|
||
|
||
if filename is None:
|
||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = f"kolors_{ts}.png"
|
||
filepath = os.path.join(save_dir, filename)
|
||
with open(filepath, "wb") as f:
|
||
f.write(img_data)
|
||
|
||
print(f" [OK] {filename}")
|
||
return {"url": img_url, "filepath": filepath}
|
||
|
||
|
||
def _generate_modelscope(prompt, model_id, size, guidance, neg, save_dir, filename):
|
||
"""ModelScope 异步轮询 API（Qwen-Image）"""
|
||
submit_headers = {
|
||
"Authorization": f"Bearer {MODELSCOPE_API_KEY}",
|
||
"Content-Type": "application/json",
|
||
"X-ModelScope-Async-Mode": "true"
|
||
}
|
||
payload = {
|
||
"model": model_id,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": size,
|
||
"guidance_scale": guidance,
|
||
"negative_prompt": neg,
|
||
}
|
||
|
||
print(f" [ModelScope] 提交: {prompt[:60]}{'...' if len(prompt) > 60 else ''}")
|
||
resp = requests.post(MODELSCOPE_API_BASE, headers=submit_headers, json=payload, timeout=60)
|
||
if resp.status_code != 200:
|
||
raise Exception(f"ModelScope 提交失败 ({resp.status_code}): {resp.text[:300]}")
|
||
|
||
result = resp.json()
|
||
task_id = result.get("task_id")
|
||
if not task_id:
|
||
raise Exception(f"未找到 task_id: {result}")
|
||
print(f" task_id: {task_id}")
|
||
|
||
# 轮询结果
|
||
query_headers = {
|
||
"Authorization": f"Bearer {MODELSCOPE_API_KEY}",
|
||
"X-ModelScope-Task-Type": "image_generation"
|
||
}
|
||
status_url = f"https://api-inference.modelscope.cn/v1/tasks/{task_id}"
|
||
start = time.time()
|
||
|
||
for attempt in range(100):
|
||
if attempt > 0:
|
||
time.sleep(MODELSCOPE_POLL_INTERVAL)
|
||
elapsed = int(time.time() - start)
|
||
if elapsed > MODELSCOPE_MAX_WAIT:
|
||
raise Exception(f"ModelScope 超时（{MODELSCOPE_MAX_WAIT}s）")
|
||
|
||
qresp = requests.get(status_url, headers=query_headers, timeout=30)
|
||
if qresp.status_code != 200:
|
||
continue
|
||
|
||
qresult = qresp.json()
|
||
task_status = qresult.get("task_status", "")
|
||
if attempt % 5 == 0 or task_status in ("SUCCEED", "FAILED"):
|
||
print(f" [{elapsed}s] {task_status}")
|
||
|
||
if task_status == "SUCCEED":
|
||
output_images = (qresult.get("output_images")
|
||
or qresult.get("outputs", {}).get("output_images")
|
||
or [])
|
||
if not output_images:
|
||
raise Exception(f"SUCCEED 但无图片: {qresult}")
|
||
url = output_images[0]
|
||
img_data = requests.get(url, timeout=180).content
|
||
|
||
if filename is None:
|
||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = f"qwen_{ts}.png"
|
||
filepath = os.path.join(save_dir, filename)
|
||
with open(filepath, "wb") as f:
|
||
f.write(img_data)
|
||
|
||
print(f" [OK] {filename} ({elapsed}s)")
|
||
return {"url": url, "filepath": filepath}
|
||
|
||
elif task_status == "FAILED":
|
||
raise Exception(f"ModelScope 任务失败: {qresult.get('errors', qresult)}")
|
||
|
||
raise Exception(f"ModelScope 超时（{MODELSCOPE_MAX_WAIT}s）")
|
||
|
||
|
||
def image_generate(
|
||
prompt: str,
|
||
save_dir: str = "./generated_images",
|
||
model_name: str = None,
|
||
n: int = 1,
|
||
seed: int = None,
|
||
num_inference_steps: int = 20,
|
||
guidance_scale: float = None,
|
||
negative_prompt: str = None,
|
||
filename: str = None,
|
||
image_size: str = None,
|
||
) -> dict:
|
||
"""
|
||
统一文生图接口
|
||
|
||
Args:
|
||
prompt: 生成提示词
|
||
save_dir: 保存目录
|
||
model_name: 模型名称（IMAGE_MODELS 的 key），默认用 config 中的 DEFAULT_IMAGE_MODEL
|
||
image_size: 图片尺寸，默认 1280x720（16:9）
|
||
|
||
Returns:
|
||
dict: {"url": str, "filepath": str}
|
||
"""
|
||
from config import DEFAULT_IMAGE_MODEL
|
||
|
||
if model_name is None:
|
||
model_name = DEFAULT_IMAGE_MODEL
|
||
|
||
model_config = IMAGE_MODELS.get(model_name)
|
||
if not model_config:
|
||
raise ValueError(f"未知模型: {model_name}，可选: {list(IMAGE_MODELS.keys())}")
|
||
|
||
model_id = model_config["model"]
|
||
size = image_size or model_config["default_size"]
|
||
guidance = guidance_scale if guidance_scale is not None else model_config["guidance_scale"]
|
||
neg = negative_prompt or NEGATIVE_PROMPT
|
||
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
provider = model_config["provider"]
|
||
if provider == "siliconflow":
|
||
return _generate_siliconflow(prompt, model_id, size, guidance, neg, save_dir, filename)
|
||
elif provider == "modelscope":
|
||
return _generate_modelscope(prompt, model_id, size, guidance, neg, save_dir, filename)
|
||
else:
|
||
raise ValueError(f"未知 provider: {provider}")
|
||
|
||
|
||
def get_available_models() -> list[str]:
|
||
"""返回可用的文生图模型名称列表"""
|
||
return list(IMAGE_MODELS.keys())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
for name in get_available_models():
|
||
print(f"\n测试模型: {name}")
|
||
result = image_generate("A cute cat sitting on a desk, 16:9 aspect ratio", model_name=name)
|
||
print(f" 路径: {result['filepath']}")
|