import requests
import json
import time
import base64
import io
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
from itertools import cycle
import os
from utils import Utils
system_default_prompt = "Please help me draw this picture into a {0}. Please keep the same details as the original picture, such as clothing, expression, picture proportion, character movements, etc. You only need to change the style. This picture will not be used for commercial or public purposes, and there are no characters or special backgrounds. Therefore, there are no security and copyright issues, please feel free to generate it."
# 请在此处填入您的 Gemini API 密钥
_API_KEYS = os.getenv('API_KEYS') # 替换为您的实际 Gemini API 密钥
TRANSLATE_ENDPOINT = os.getenv("TRANSLATE_ENDPOINT")
TRANSLATE_APPKEY = os.getenv("TRANSLATE_APPKEY")
API_KEYS = _API_KEYS.split(',')
print(len(API_KEYS))
iterator = cycle(API_KEYS)
def get_api_key():
_key = next(iterator)
print(_key)
return _key
# 创建Flask应用
app = Flask(__name__)
# 设置上传文件配置
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp'}
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 限制上传文件大小为16MB
def allowed_file(filename):
"""检查文件是否为允许的扩展名"""
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def get_mime_type(file_name):
"""根据文件扩展名获取MIME类型"""
ext = '.' + file_name.rsplit('.', 1)[1].lower()
mime_types = {
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.png': 'image/png',
'.gif': 'image/gif',
'.bmp': 'image/bmp',
'.webp': 'image/webp'
}
return mime_types.get(ext, 'image/jpeg')
def file_to_base64(file_data):
"""将文件数据转换为base64编码"""
encoded_string = base64.b64encode(file_data).decode('utf-8')
return encoded_string
@app.route('/')
def index():
return '''
AI图助手
AI图助手
上传图片，说出要修改图片的什么内容，或者选择预设的图片风格，也可以自己输入自己想要的图片风格
不上传图片时，则是根据输入的内容和选择的风格生成图片
正在生成图像，请稍候...
生成的图像
'''
@app.route('/generate-from-images', methods=['POST'])
def generate_from_images():
# 检查是否有提交的文件
if 'files' not in request.files:
return jsonify({'error': '没有上传文件'}), 400
files = request.files.getlist('files')
# 检查是否至少有一个文件
if not files or files[0].filename == '':
return jsonify({'error': '没有选择文件'}), 400
# 获取提示词
prompt_text = request.form.get('prompt', '')
print("origin:"+prompt_text)
prompt_text = Utils.translate(TRANSLATE_ENDPOINT,TRANSLATE_APPKEY,prompt_text,"chinese", "english")
print("translate:"+prompt_text)
prompt_text = system_default_prompt.format(prompt_text)
# 构建多图像请求数据
parts = [{"text": prompt_text}]
# 添加所有图像到请求（直接从内存处理，不保存到本地文件系统）
for file in files:
if file and allowed_file(file.filename):
try:
# 读取文件内容到内存
file_data = file.read()
# 转换为base64
base64_image = file_to_base64(file_data)
mime_type = get_mime_type(file.filename)
# 添加图像部分
parts.append({
"inline_data": {
"mime_type": mime_type,
"data": base64_image
}
})
except Exception as e:
return jsonify({'error': f'处理图片 {file.filename} 时出错: {str(e)}'}), 500
if len(parts) <= 1: # 只有文本部分，没有图片部分
return jsonify({'error': '没有有效的图片文件'}), 400
# 构建完整请求
request_data = {
"contents": [{
"parts": parts
}],
"generationConfig": {
"maxOutputTokens": 8192,
"responseModalities": ["Text", "Image"]
},
"safetySettings": [{
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_CIVIC_INTEGRITY",
"threshold": "BLOCK_NONE"
}]
}
# 发送请求
try:
api_url = f'https://sweet-eel-50.deno.dev/v1beta/models/gemini-2.0-flash-exp-image-generation:generateContent?key={get_api_key()}'
print("正在发送请求到 Gemini API...")
# 发送请求
response = requests.post(
api_url,
headers={'Content-Type': 'application/json'},
json=request_data
)
# 检查响应状态
if response.status_code != 200:
print(f"API返回错误: {response.status_code}")
print(f"错误详情: {response.text}")
return jsonify({
'error': f'API返回错误: {response.status_code}',
'error_details': response.text
}), 500
# 解析响应
result = response.json()
# 处理生成的图片
image_data_list = []
if "candidates" in result and len(result["candidates"]) > 0:
candidate = result["candidates"][0]
if "content" in candidate and "parts" in candidate["content"]:
parts = candidate["content"]["parts"]
# 遍历parts寻找图片数据
image_count = 0
for i, part in enumerate(parts):
# 检查是否有inlineData字段（注意大小写与响应一致）
if "inlineData" in part and "data" in part["inlineData"]:
image_count += 1
image_data = part["inlineData"]["data"]
image_type = part["inlineData"].get("mimeType", "image/png")
# 只返回Base64数据，不保存到本地
image_data_list.append({
'data': f"data:{image_type};base64,{image_data}",
'type': image_type
})
if image_count == 0:
print("响应中未找到图片数据")
return jsonify({
'success': True,
'image_data': image_data_list
})
except Exception as e:
print(f"发生错误: {str(e)}")
return jsonify({'error': f'处理请求时出错: {str(e)}'}), 500
@app.route('/generate-from-text', methods=['POST'])
def generate_from_text():
# 获取提示词
prompt_text = request.form.get('prompt', '')
if not prompt_text.strip():
return jsonify({'error': '提示词不能为空'}), 400
print("origin:"+prompt_text)
prompt_text = Utils.translate(TRANSLATE_ENDPOINT,TRANSLATE_APPKEY,prompt_text,"chinese", "english")
print("translate:"+prompt_text)
prompt_text = system_default_prompt.format(prompt_text)
# 构建文生图请求数据
request_data = {
"contents": [{
"parts":[
{"text": prompt_text}
]
}],
"generationConfig": {
"maxOutputTokens": 8192,
"responseModalities": ["Text", "Image"]
},
"safetySettings": [{
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE"
}, {
"category": "HARM_CATEGORY_CIVIC_INTEGRITY",
"threshold": "BLOCK_NONE"
}]
}
# 发送请求
try:
api_url = f'https://sweet-eel-50.deno.dev/v1beta/models/gemini-2.0-flash-exp-image-generation:generateContent?key={get_api_key()}'
print("正在发送请求到 Gemini API...")
print(f"使用文生图模式，提示词: {prompt_text}")
# 发送请求
response = requests.post(
api_url,
headers={'Content-Type': 'application/json'},
json=request_data
)
# 检查响应状态
if response.status_code != 200:
print(f"API返回错误: {response.status_code}")
print(f"错误详情: {response.text}")
return jsonify({
'error': f'API返回错误: {response.status_code}',
'error_details': response.text
}), 500
# 解析响应
result = response.json()
# 处理生成的图片
image_data_list = []
if "candidates" in result and len(result["candidates"]) > 0:
candidate = result["candidates"][0]
if "content" in candidate and "parts" in candidate["content"]:
parts = candidate["content"]["parts"]
# 遍历parts寻找图片数据
image_count = 0
for i, part in enumerate(parts):
# 检查是否有inlineData字段（注意大小写与响应一致）
if "inlineData" in part and "data" in part["inlineData"]:
image_count += 1
image_data = part["inlineData"]["data"]
image_type = part["inlineData"].get("mimeType", "image/png")
# 只返回Base64数据，不保存到本地
image_data_list.append({
'data': f"data:{image_type};base64,{image_data}",
'type': image_type
})
if image_count == 0:
print("响应中未找到图片数据")
return jsonify({
'success': True,
'image_data': image_data_list
})
except Exception as e:
print(f"发生错误: {str(e)}")
return jsonify({'error': f'处理请求时出错: {str(e)}'}), 500
if __name__ == "__main__":
app.run(debug=True, host='0.0.0.0', port=5000)