diff --git a/.runtime.env b/.runtime.env
index a31361e..0236cf0 100644
--- a/.runtime.env
+++ b/.runtime.env
@@ -1,5 +1,5 @@
-MODAL_ENVIRONMENT=test
-modal_app_name=bowong-ai-video
+MODAL_ENVIRONMENT=dev
+modal_app_name=bowong-ai-video-gemini
S3_mount_dir=/mntS3
S3_bucket_name=modal-media-cache
S3_region=ap-northeast-2
diff --git a/src/BowongModalFunctions/router/google.py b/src/BowongModalFunctions/router/google.py
index 7699cb7..70932cd 100644
--- a/src/BowongModalFunctions/router/google.py
+++ b/src/BowongModalFunctions/router/google.py
@@ -1,26 +1,26 @@
import json
import os
-from typing import Annotated, Optional, cast, List
+from typing import Annotated, Optional, cast
+import httpx
import modal
import sentry_sdk
+from fastapi import APIRouter, UploadFile, Header, HTTPException, Depends, File, Form
from fastapi.responses import Response
from loguru import logger
-import httpx
-from fastapi import APIRouter, UploadFile, Header, HTTPException, Depends
from pydantic import BaseModel, Field
from pydantic import computed_field
from starlette import status
from starlette.responses import JSONResponse
-from ..models.settings.cluster import WorkerConfig
from ..middleware.authorization import verify_token
-from ..models.responses.models import ModalTaskResponse, GeminiResultResponse, GoogleUploadResponse
+from ..models.ffmpeg_tasks.models import SentryTransactionInfo
from ..models.requests.models import GeminiRequest, MakeGridGeminiRequest, GeminiRequestFirstOnly, \
GeminiRequestSecondOnly
-from ..models.ffmpeg_tasks.models import SentryTransactionInfo
-from ..utils.ModalUtils import ModalUtils
+from ..models.responses.models import ModalTaskResponse, GeminiResultResponse, GoogleUploadResponse
+from ..models.settings.cluster import WorkerConfig
from ..utils.HTTPUtils import GoogleAuthUtils
+from ..utils.ModalUtils import ModalUtils
config = WorkerConfig()
@@ -176,23 +176,7 @@ async def get_access_token() -> GoogleAuthUtils.GoogleAuthResponse:
return access_token
-@router.post("/make_grid_gemini", summary="将输入图拼为网格上传到Gemini网盘")
-async def make_grid_gemini_upload(data: MakeGridGeminiRequest, headers: Annotated[BundleHeaders, Header()]):
- google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
- if not google_api_key:
- raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
- fn = modal.Function.from_name(config.modal_app_name, "make_image_grid_upload",
- environment_name=config.modal_environment)
- image_grid_uri = await fn.remote.aio(data.pic_info_list, data.image_size, data.text_height, data.font_size,
- data.padding, data.separator, google_api_key,
- SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
- x_baggage=sentry_sdk.get_baggage())
- if headers.x_trace_id is None else SentryTransactionInfo(
- x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage))
- return JSONResponse(content={"uri": image_grid_uri}, status_code=200)
-
-
-@router.post('/inference_gemini', summary="使用Gemini推理hls视频流指定时间段的打点情况")
+@router.post('/video/inference_gemini', summary="使用Gemini推理hls视频流指定时间段的打点情况")
async def inference_gemini(data: GeminiRequest, headers: Annotated[BundleHeaders, Header()], ) -> ModalTaskResponse:
google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
if not google_api_key:
@@ -209,7 +193,7 @@ async def inference_gemini(data: GeminiRequest, headers: Annotated[BundleHeaders
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
-@router.post('/inference_gemini_first_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第一阶段")
+@router.post('/video/inference_gemini_first_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第一阶段")
async def inference_gemini_first_only(data: GeminiRequestFirstOnly,
headers: Annotated[BundleHeaders, Header()], ) -> ModalTaskResponse:
google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
@@ -221,7 +205,7 @@ async def inference_gemini_first_only(data: GeminiRequestFirstOnly,
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
-@router.post('/inference_gemini_second_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第二阶段")
+@router.post('/video/inference_gemini_second_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第二阶段")
async def inference_gemini_second_only(data: GeminiRequestSecondOnly,
headers: Annotated[BundleHeaders, Header()], ) -> ModalTaskResponse:
google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
@@ -234,7 +218,54 @@ async def inference_gemini_second_only(data: GeminiRequestSecondOnly,
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
-@router.get("/inference_gemini/{task_id}", summary="查询Gemini推理hls视频流指定时间段的打点任务")
+@router.post("/image/make_grid_gemini", summary="将输入图拼为网格上传到Gemini网盘")
+async def make_grid_gemini_upload(data: MakeGridGeminiRequest, headers: Annotated[BundleHeaders, Header()]):
+ google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
+ if not google_api_key:
+ raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
+ fn = modal.Function.from_name(config.modal_app_name, "make_image_grid_upload",
+ environment_name=config.modal_environment)
+ image_grid_uri = await fn.remote.aio(data.pic_info_list, data.image_size, data.text_height, data.font_size,
+ data.padding, data.separator, google_api_key,
+ SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
+ x_baggage=sentry_sdk.get_baggage())
+ if headers.x_trace_id is None else SentryTransactionInfo(
+ x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage))
+ return JSONResponse(content={"uri": image_grid_uri}, status_code=200)
+
+
+@router.post("/image/clothes_mark", summary="使用Gemini将图像中模特身上某一件衣服涂成指定颜色")
+async def clothes_mark(
+ headers: Annotated[BundleHeaders, Header()],
+ origin_image: Annotated[UploadFile, File(description="待处理图片")],
+ mark_clothes_type: Annotated[str, Form(description="替换衣服类型")],
+ mark_color: Annotated[str, Form(description="替换颜色")] = "绿色"
+) -> ModalTaskResponse:
+ google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
+ if not google_api_key:
+ raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
+ fn = modal.Function.from_name(config.modal_app_name, "image_clothes_mark",
+ environment_name=config.modal_environment)
+ fn_call = fn.spawn(google_api_key, origin_image, mark_clothes_type, mark_color)
+ return ModalTaskResponse(success=True, taskId=fn_call.object_id)
+
+
+@router.post("/image/edit_custom", summary="使用Gemini进行图像编辑")
+async def clothes_mark(
+ headers: Annotated[BundleHeaders, Header()],
+ origin_image: Annotated[UploadFile, File(description="待处理图片")],
+ prompt: Annotated[str, Form(description="图像编辑提示词")]
+) -> ModalTaskResponse:
+ google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
+ if not google_api_key:
+ raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
+ fn = modal.Function.from_name(config.modal_app_name, "image_edit_custom",
+ environment_name=config.modal_environment)
+ fn_call = fn.spawn(google_api_key, origin_image, prompt)
+ return ModalTaskResponse(success=True, taskId=fn_call.object_id)
+
+
+@router.get("/{task_id}", summary="查询Gemini推理结果")
async def gemini_status(task_id: str, response: Response) -> GeminiResultResponse:
task_info = await ModalUtils.get_modal_task_status(task_id)
if task_info.transaction:
diff --git a/src/BowongModalFunctions/utils/HTTPUtils.py b/src/BowongModalFunctions/utils/HTTPUtils.py
index 3c150c4..421f012 100644
--- a/src/BowongModalFunctions/utils/HTTPUtils.py
+++ b/src/BowongModalFunctions/utils/HTTPUtils.py
@@ -155,14 +155,14 @@ class GoogleAuthUtils:
regions: List[str]
access_token: str
- @computed_field(description="通过Cloudflare gateway调用的url")
+ @computed_field(description="通过谷歌官方API调用的url")
@property
def gateway_url(self) -> str:
- return f"https://gateway.ai.cloudflare.com/v1/{self.cloudflare_project_id}/{self.cloudflare_gateway_id}/google-vertex-ai/v1/projects/{self.google_project_id}/locations/{random.choice(self.regions)}/publishers/google/models"
+ return f"https://aiplatform.googleapis.com/v1beta1/projects/{self.google_project_id}/locations/{random.choice(self.regions)}/publishers/google/models"
def generate_content(self, model_id: str, contents: Union[types.ContentListUnion, types.ContentListUnionDict],
config: Optional[types.GenerateContentConfig] = None,
- timeout: int = 30) -> tuple[dict[Any, Any], int] | tuple[GenerateContentResponse, int]:
+ timeout: int = 30) -> tuple[str, int] | tuple[GenerateContentResponse, int]:
parameter_model = GoogleAuthUtils.VertexAIRequestModel(contents=contents, generationConfig=config)
json_body = parameter_model.model_dump_json(indent=2, exclude_none=True, by_alias=True,
exclude={"generation_config": ["safety_settings", "system_instruction"]})
@@ -177,7 +177,7 @@ class GoogleAuthUtils:
# response.raise_for_status()
# logger.info(response.text)
if response.status_code != 200:
- return {}, response.status_code
+ return response.text, response.status_code
return types.GenerateContentResponse.model_validate_json(response.text), response.status_code
@staticmethod
diff --git a/src/cluster/image_apps/__init__.py b/src/cluster/image_apps/__init__.py
new file mode 100644
index 0000000..807d170
--- /dev/null
+++ b/src/cluster/image_apps/__init__.py
@@ -0,0 +1,17 @@
+import os
+import importlib
+import pkgutil
+
+# 获取当前包所在的目录
+package_dir = os.path.dirname(__file__)
+
+# 自动导入所有模块
+for (_, module_name, _) in pkgutil.iter_modules([package_dir]):
+ # 导入模块
+ module = importlib.import_module(f'.{module_name}', package=__name__)
+ # 将模块添加到当前包的命名空间
+ globals()[module_name] = module
+ # 可选：将模块名添加到 __all__
+ if '__all__' not in globals():
+ globals()['__all__'] = []
+ globals()['__all__'].append(module_name)
\ No newline at end of file
diff --git a/src/cluster/image_apps/image_edit.py b/src/cluster/image_apps/image_edit.py
new file mode 100644
index 0000000..8674c3a
--- /dev/null
+++ b/src/cluster/image_apps/image_edit.py
@@ -0,0 +1,241 @@
+import base64
+
+import modal
+
+from ..video import app, config, downloader_image
+
+with (downloader_image.imports()):
+ import json
+ from google.genai import types
+ from loguru import logger
+ from fastapi import UploadFile, HTTPException
+ import tempfile
+ import uuid
+ from BowongModalFunctions.utils.HTTPUtils import GoogleAuthUtils
+
+ gemini_region = [
+ "global"
+ ]
+
+ SUPPORTED_IMAGE_TYPES = {
+ "image/jpeg": ".jpg",
+ "image/png": ".png",
+ "image/webp": ".webp"
+ }
+
+
+ async def upload(file_path: str, content_type: str,
+ google_api_key: str) -> str:
+ """
+ 上传视频到gcs
+ """
+ with open(file_path, "rb") as image:
+ image = image.read()
+ content_length = len(image)
+ filename = ".".join([str(uuid.uuid4()), SUPPORTED_IMAGE_TYPES[content_type][1:]])
+ logger.info(
+ f"Uploading name = {filename}, size = {content_length}, type = {content_type} to google file")
+ upload_response = await GoogleAuthUtils.google_upload_file(file_stream=image,
+ content_type=content_type,
+ google_api_key=google_api_key,
+ bucket_name="dy-media-storage",
+ filename=f"images/{filename}", )
+
+ return upload_response.urn
+
+
+ @app.function(cpu=(0.5, 64), timeout=1800,
+ max_containers=config.video_downloader_concurrency,
+ volumes={
+ config.S3_mount_dir: modal.CloudBucketMount(
+ bucket_name=config.S3_bucket_name,
+ secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.modal_environment),
+ ),
+ },
+ region="us",
+ cloud='gcp')
+ @modal.concurrent(max_inputs=1)
+ async def image_clothes_mark(google_api_key: str, origin_image: UploadFile, mark_clothes_type: str, mark_color: str):
+ content_type = origin_image.content_type
+ if content_type not in SUPPORTED_IMAGE_TYPES:
+ raise HTTPException(
+ status_code=400,
+ detail=f"不支持的文件类型: {content_type}。支持的类型: {list(SUPPORTED_IMAGE_TYPES.keys())}"
+ )
+
+ try:
+ # 获取对应文件类型的扩展名
+ file_extension = SUPPORTED_IMAGE_TYPES[content_type]
+ # 创建临时文件，使用正确的扩展名
+ with tempfile.NamedTemporaryFile(
+ delete=False,
+ suffix=file_extension,
+ mode='wb'
+ ) as temp_file:
+ # 写入文件内容
+ contents = await origin_image.read()
+ temp_file.write(contents)
+ temp_file_path = temp_file.name
+ logger.success(f"上传图片已保存到{temp_file_path}")
+ except Exception as e:
+ logger.exception("保存图片失败")
+ raise e
+ origin_image_uri = await upload(temp_file_path, content_type, google_api_key)
+ client = GoogleAuthUtils.GoogleGenaiClient(
+ cloudflare_project_id="67720b647ff2b55cf37ba3ef9e677083",
+ cloudflare_gateway_id="bowong-dev",
+ google_project_id="gen-lang-client-0413414134",
+ regions=gemini_region, access_token=google_api_key,
+ )
+ try:
+ logger.info("🐬开始处理图片")
+ resp, resp_code = client.generate_content(model_id="gemini-2.0-flash-preview-image-generation",
+ contents=[types.Content(role='user',
+ parts=[
+ types.Part.from_uri(
+ file_uri=origin_image_uri,
+ mime_type=content_type),
+ types.Part.from_text(
+ text=""
+ ""
+ "请将图片中模特身上的**{0}**涂为**纯{1}**, 不考虑褶皱纹理透视，用于服装替换, 其他部分保持不变, 输出原尺寸的图片"
+ ""
+ "".format(
+ mark_clothes_type, mark_color
+ )
+ )
+ ]
+ )
+ ],
+ config=types.GenerateContentConfig.model_validate({
+ "temperature": 0.01,
+ "top_p": 0.7,
+ "response_modalities": ["IMAGE","TEXT"]
+ }),
+ timeout=900)
+ except Exception as e:
+ logger.exception(f"😭修改图片请求失败 {e}")
+ return None, None
+ if resp_code == 200:
+ reason = resp.candidates[0].finish_reason
+ logger.info("😊用量信息" + resp.usage_metadata.model_dump_json(indent=2))
+ if reason == "STOP" or reason == "FINISH_REASON_UNSPECIFIED":
+ for part in resp.candidates[0].content.parts:
+ if part.inline_data is not None:
+ result_image_raw: bytes = part.inline_data.data
+ result_image_mime: str = part.inline_data.mime_type
+ result_image_b64 = base64.b64encode(result_image_raw).decode('ascii')
+ result = {
+ "image_b64": f"data:{result_image_mime};base64,{result_image_b64}",
+ }
+ # 解析识别结果
+ logger.info(f"Result:\n {json.dumps(result, indent=2)}")
+ return result, None
+ else:
+ logger.error(f"😭Gemini推理失败, Reason {reason}")
+ return None, None
+ else:
+ logger.error(f"😭Gemini推理失败, 状态码{resp_code}")
+ if resp:
+ logger.error(f"😭Gemini推理失败, 返回内容{resp}")
+ if resp_code == 429:
+ logger.warning("🥵请求负载过高, 随机更换地区重试")
+ return None, None
+
+ @app.function(cpu=(0.5, 64), timeout=1800,
+ max_containers=config.video_downloader_concurrency,
+ volumes={
+ config.S3_mount_dir: modal.CloudBucketMount(
+ bucket_name=config.S3_bucket_name,
+ secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.modal_environment),
+ ),
+ },
+ region="us",
+ cloud='gcp')
+ @modal.concurrent(max_inputs=1)
+ async def image_edit_custom(google_api_key: str, origin_image: UploadFile, prompt: str):
+ content_type = origin_image.content_type
+ if content_type not in SUPPORTED_IMAGE_TYPES:
+ raise HTTPException(
+ status_code=400,
+ detail=f"不支持的文件类型: {content_type}。支持的类型: {list(SUPPORTED_IMAGE_TYPES.keys())}"
+ )
+
+ try:
+ # 获取对应文件类型的扩展名
+ file_extension = SUPPORTED_IMAGE_TYPES[content_type]
+ # 创建临时文件，使用正确的扩展名
+ with tempfile.NamedTemporaryFile(
+ delete=False,
+ suffix=file_extension,
+ mode='wb'
+ ) as temp_file:
+ # 写入文件内容
+ contents = await origin_image.read()
+ temp_file.write(contents)
+ temp_file_path = temp_file.name
+ logger.success(f"上传图片已保存到{temp_file_path}")
+ except Exception as e:
+ logger.exception("保存图片失败")
+ raise e
+ origin_image_uri = await upload(temp_file_path, content_type, google_api_key)
+ client = GoogleAuthUtils.GoogleGenaiClient(
+ cloudflare_project_id="67720b647ff2b55cf37ba3ef9e677083",
+ cloudflare_gateway_id="bowong-dev",
+ google_project_id="gen-lang-client-0413414134",
+ regions=gemini_region, access_token=google_api_key,
+ )
+ try:
+ logger.info("🐬开始处理图片")
+ resp, resp_code = client.generate_content(model_id="gemini-2.0-flash-preview-image-generation",
+ contents=[types.Content(role='user',
+ parts=[
+ types.Part.from_uri(
+ file_uri=origin_image_uri,
+ mime_type=content_type),
+ types.Part.from_text(
+ text=""
+ ""
+ "{0}"
+ ""
+ "".format(
+ prompt
+ )
+ )
+ ]
+ )
+ ],
+ config=types.GenerateContentConfig.model_validate({
+ "temperature": 0.01,
+ "top_p": 0.7,
+ "response_modalities": ["IMAGE","TEXT"]
+ }),
+ timeout=900)
+ except Exception as e:
+ logger.exception(f"😭修改图片请求失败 {e}")
+ return None, None
+ if resp_code == 200:
+ reason = resp.candidates[0].finish_reason
+ logger.info("😊用量信息" + resp.usage_metadata.model_dump_json(indent=2))
+ if reason == "STOP" or reason == "FINISH_REASON_UNSPECIFIED":
+ for part in resp.candidates[0].content.parts:
+ if part.inline_data is not None:
+ result_image_raw: bytes = part.inline_data.data
+ result_image_mime: str = part.inline_data.mime_type
+ result_image_b64 = base64.b64encode(result_image_raw).decode('ascii')
+ result = {
+ "image_b64": f"data:{result_image_mime};base64,{result_image_b64}",
+ }
+ # 解析识别结果
+ logger.info(f"Result:\n {json.dumps(result, indent=2)}")
+ return result, None
+ else:
+ logger.error(f"😭Gemini推理失败, Reason {reason}")
+ return None, None
+ else:
+ logger.error(f"😭Gemini推理失败, 状态码{resp_code}")
+ if resp:
+ logger.error(f"😭Gemini推理失败, 返回内容{resp}")
+ if resp_code == 429:
+ logger.warning("🥵请求负载过高, 随机更换地区重试")
+ return None, None
diff --git a/src/cluster/video.py b/src/cluster/video.py
index 06b70c7..c35a651 100644
--- a/src/cluster/video.py
+++ b/src/cluster/video.py
@@ -106,6 +106,7 @@ with downloader_image.imports():
# 导入sub app function
from .video_apps import *
+ from .image_apps import *