ADD 增加图像编辑推理接口
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
17
src/cluster/image_apps/__init__.py
Normal file
17
src/cluster/image_apps/__init__.py
Normal 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)
|
||||
241
src/cluster/image_apps/image_edit.py
Normal file
241
src/cluster/image_apps/image_edit.py
Normal 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
|
||||
@@ -106,6 +106,7 @@ with downloader_image.imports():
|
||||
|
||||
# 导入sub app function
|
||||
from .video_apps import *
|
||||
from .image_apps import *
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user