mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-06-05 15:49:38 +08:00
a74ac03836
- 添加图片生成相关配置和环境变量 - 新增图片上传服务和模型定义 - 扩展模型服务以支持图片生成模型 - 添加图片生成响应处理器 - 更新README文档以反映新功能 - 添加GitHub Actions发布工作流
81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
import time
|
|
import uuid
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
import base64
|
|
|
|
from app.core.config import settings
|
|
from app.core.logger import get_image_create_logger
|
|
from app.core.uploader import ImageUploaderFactory
|
|
from app.schemas.openai_models import ImageGenerationRequest
|
|
|
|
logger = get_image_create_logger()
|
|
|
|
|
|
class ImageCreateService:
|
|
def __init__(self, aspect_ratio="1:1"):
|
|
self.image_model = settings.CREATE_IMAGE_MODEL
|
|
self.paid_key = settings.PAID_KEY
|
|
self.aspect_ratio = aspect_ratio
|
|
|
|
def generate_images(self, request: ImageGenerationRequest):
|
|
client = genai.Client(api_key=self.paid_key)
|
|
if request.size == "1024x1024":
|
|
self.aspect_ratio = "1:1"
|
|
elif request.size == "1792x1024":
|
|
self.aspect_ratio = "16:9"
|
|
elif request.size == "1027x1792":
|
|
self.aspect_ratio = "9:16"
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid size: {request.size}. Supported sizes are 1024x1024, 1792x1024, and 1024x1792."
|
|
)
|
|
|
|
response = client.models.generate_images(
|
|
model=self.image_model,
|
|
prompt=request.prompt,
|
|
config=types.GenerateImagesConfig(
|
|
number_of_images=request.n,
|
|
output_mime_type="image/png",
|
|
aspect_ratio=self.aspect_ratio,
|
|
safety_filter_level="BLOCK_LOW_AND_ABOVE",
|
|
person_generation="ALLOW_ADULT",
|
|
# language="auto"
|
|
),
|
|
)
|
|
|
|
if response.generated_images:
|
|
images_data = []
|
|
for index, generated_image in enumerate(response.generated_images):
|
|
image_data = generated_image.image.image_bytes
|
|
image_uploader = None
|
|
if settings.UPLOAD_PROVIDER == "smms":
|
|
image_uploader = ImageUploaderFactory.create(provider=settings.UPLOAD_PROVIDER,api_key=settings.SMMS_SECRET_TOKEN)
|
|
current_date = time.strftime("%Y/%m/%d")
|
|
filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
|
|
upload_response = image_uploader.upload(image_data,filename)
|
|
|
|
# base64_image = base64.b64encode(image_data).decode('utf-8')
|
|
images_data.append({
|
|
"url": f"{upload_response.data.url}",
|
|
"revised_prompt": request.prompt
|
|
})
|
|
|
|
response_data = {
|
|
"created": int(time.time()), # Current timestamp
|
|
"data": images_data
|
|
}
|
|
return response_data
|
|
else:
|
|
raise Exception("I can't generate these images")
|
|
|
|
def generate_images_chat(self, request: ImageGenerationRequest) -> str:
|
|
response = self.generate_images(request)
|
|
image_datas = response["data"]
|
|
if image_datas:
|
|
markdown_images = []
|
|
for index, image_data in enumerate(image_datas):
|
|
markdown_images.append(f"")
|
|
return "\n".join(markdown_images)
|
|
|