feat: 添加图片生成功能及相关配置

- 添加图片生成相关配置和环境变量 - 新增图片上传服务和模型定义 - 扩展模型服务以支持图片生成模型 - 添加图片生成响应处理器 - 更新README文档以反映新功能 - 添加GitHub Actions发布工作流
2026-09-04 23:19:17 +08:00 · 2025-02-11 01:59:16 +08:00
parent b3842b2329
commit a74ac03836
16 changed files with 497 additions and 28 deletions
@@ -0,0 +1,81 @@
+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"![Generated Image {index+1}]({image_data['url']})")
+            return "\n".join(markdown_images)
+