|
from flask import Flask, request, jsonify, Response |
|
import requests |
|
import json |
|
import time |
|
import random |
|
import logging |
|
from requests.exceptions import RequestException |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') |
|
logger = logging.getLogger(__name__) |
|
|
|
SYSTEM_ASSISTANT = """作为 Stable Diffusion Prompt 提示词专家，您将从关键词中创建提示，通常来自 Danbooru 等数据库。 |
|
提示通常描述图像，使用常见词汇，按重要性排列，并用逗号分隔。避免使用"-"或"."，但可以接受空格和自然语言。避免词汇重复。 |
|
|
|
为了强调关键词，请将其放在括号中以增加其权重。例如，"(flowers)"将'flowers'的权重增加1.1倍，而"(((flowers)))"将其增加1.331倍。使用"(flowers:1.5)"将'flowers'的权重增加1.5倍。只为重要的标签增加权重。 |
|
|
|
提示包括三个部分：**前缀** （质量标签+风格词+效果器）+ **主题** （图像的主要焦点）+ **场景** （背景、环境）。 |
|
|
|
* 前缀影响图像质量。像"masterpiece"、"best quality"、"4k"这样的标签可以提高图像的细节。像"illustration"、"lensflare"这样的风格词定义图像的风格。像"bestlighting"、"lensflare"、"depthoffield"这样的效果器会影响光照和深度。 |
|
|
|
* 主题是图像的主要焦点，如角色或场景。对主题进行详细描述可以确保图像丰富而详细。增加主题的权重以增强其清晰度。对于角色，描述面部、头发、身体、服装、姿势等特征。 |
|
|
|
* 场景描述环境。没有场景，图像的背景是平淡的，主题显得过大。某些主题本身包含场景（例如建筑物、风景）。像"花草草地"、"阳光"、"河流"这样的环境词可以丰富场景。你的任务是设计图像生成的提示。请按照以下步骤进行操作： |
|
|
|
1. 我会发送给您一个图像场景。需要你生成详细的图像描述 |
|
2. 图像描述必须是英文，输出为Positive Prompt。 |
|
|
|
示例： |
|
|
|
我发送：二战时期的护士。 |
|
您回复只回复： |
|
A WWII-era nurse in a German uniform, holding a wine bottle and stethoscope, sitting at a table in white attire, with a table in the background, masterpiece, best quality, 4k, illustration style, best lighting, depth of field, detailed character, detailed environment. |
|
""" |
|
|
|
def get_random_token(auth_header): |
|
if not auth_header: |
|
return None |
|
if auth_header.startswith('Bearer '): |
|
auth_header = auth_header[7:] |
|
tokens = [token.strip() for token in auth_header.split(',') if token.strip()] |
|
if not tokens: |
|
return None |
|
return f"Bearer {random.choice(tokens)}" |
|
|
|
def translate_and_enhance_prompt(prompt, auth_token): |
|
translate_url = 'https://api.siliconflow.cn/v1/chat/completions' |
|
translate_body = { |
|
'model': 'Qwen/Qwen2-72B-Instruct', |
|
'messages': [ |
|
{'role': 'system', 'content': SYSTEM_ASSISTANT}, |
|
{'role': 'user', 'content': prompt} |
|
] |
|
} |
|
headers = { |
|
'Content-Type': 'application/json', |
|
'Authorization': auth_token |
|
} |
|
try: |
|
response = requests.post(translate_url, headers=headers, json=translate_body, timeout=30) |
|
response.raise_for_status() |
|
result = response.json() |
|
return result['choices'][0]['message']['content'] |
|
except RequestException as e: |
|
logger.error(f"Error in translate_and_enhance_prompt: {str(e)}") |
|
raise |
|
|
|
@app.route('/') |
|
def index(): |
|
return "text-to-image with siliconflow", 200 |
|
|
|
@app.route('/ai/v1/chat/completions', methods=['POST']) |
|
def handle_request(): |
|
try: |
|
body = request.json |
|
model = body.get('model') |
|
messages = body.get('messages') |
|
stream = body.get('stream', False) |
|
if not model or not messages or len(messages) == 0: |
|
return jsonify({"error": "Bad Request: Missing required fields"}), 400 |
|
|
|
prompt = messages[-1]['content'] |
|
|
|
random_token = get_random_token(request.headers.get('Authorization')) |
|
if not random_token: |
|
return jsonify({"error": "Unauthorized: Invalid or missing Authorization header"}), 401 |
|
|
|
try: |
|
enhanced_prompt = translate_and_enhance_prompt(prompt, random_token) |
|
except Exception as e: |
|
logger.error(f"Error in translate_and_enhance_prompt: {str(e)}") |
|
return jsonify({"error": "Failed to enhance prompt"}), 500 |
|
|
|
new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image' |
|
new_request_body = { |
|
"prompt": enhanced_prompt, |
|
"image_size": "1024x1024", |
|
"batch_size": 1, |
|
"num_inference_steps": 4, |
|
"guidance_scale": 1 |
|
} |
|
|
|
headers = { |
|
'accept': 'application/json', |
|
'content-type': 'application/json', |
|
'Authorization': random_token |
|
} |
|
|
|
try: |
|
response = requests.post(new_url, headers=headers, json=new_request_body, timeout=60) |
|
response.raise_for_status() |
|
response_body = response.json() |
|
image_url = response_body['images'][0]['url'] |
|
except RequestException as e: |
|
logger.error(f"Error in image generation request: {str(e)}") |
|
return jsonify({"error": "Failed to generate image"}), 500 |
|
|
|
unique_id = int(time.time() * 1000) |
|
current_timestamp = unique_id // 1000 |
|
|
|
if stream: |
|
response_payload = { |
|
"id": unique_id, |
|
"object": "chat.completion.chunk", |
|
"created": current_timestamp, |
|
"model": model, |
|
"choices": [ |
|
{ |
|
"index": 0, |
|
"delta": { |
|
"content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n" |
|
}, |
|
"finish_reason": "stop" |
|
} |
|
] |
|
} |
|
data_string = json.dumps(response_payload) |
|
return Response(f"data: {data_string}\n\n", content_type='text/event-stream') |
|
else: |
|
response_payload = { |
|
"id": unique_id, |
|
"object": "chat.completion", |
|
"created": current_timestamp, |
|
"model": model, |
|
"choices": [ |
|
{ |
|
"index": 0, |
|
"message": { |
|
"role": "assistant", |
|
"content": f"原始提示词：{prompt}\n增强后的提示词：{enhanced_prompt}\n生成的图像：\n" |
|
}, |
|
"logprobs": None, |
|
"finish_reason": "length" |
|
} |
|
], |
|
"usage": { |
|
"prompt_tokens": len(prompt), |
|
"completion_tokens": len(enhanced_prompt) + len(image_url), |
|
"total_tokens": len(prompt) + len(enhanced_prompt) + len(image_url) |
|
} |
|
} |
|
data_string = json.dumps(response_payload) |
|
return Response(f"{data_string}\n\n", content_type='text/event-stream') |
|
except Exception as e: |
|
logger.error(f"Unexpected error in handle_request: {str(e)}") |
|
return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500 |
|
|
|
if __name__ == '__main__': |
|
app.run(host='0.0.0.0', port=8000) |
|
|
