ADD 增加图像编辑推理接口

This commit is contained in:
2025-07-11 15:21:00 +08:00
parent 720d8e212c
commit 4c6c49b7ba
6 changed files with 323 additions and 33 deletions

View File

@@ -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

View File

@@ -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:

View File

@@ -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

View File

@@ -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)

View File

@@ -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="<prompt>"
"<instruction>"
"请将图片中模特身上的**{0}**涂为**纯{1}**, 不考虑褶皱纹理透视,用于服装替换, 其他部分保持不变, 输出原尺寸的图片"
"</instruction>"
"</prompt>".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="<prompt>"
"<instruction>"
"{0}"
"</instruction>"
"</prompt>".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

View File

@@ -106,6 +106,7 @@ with downloader_image.imports():
# 导入sub app function
from .video_apps import *
from .image_apps import *