12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273 |
- import uvicorn
- import asyncio
- from fastapi import FastAPI
- from pydantic import BaseModel
- from google_ai.generativeai_video import main
- app = FastAPI()
- api_keys = [
- {'key': 'AIzaSyB2kjF2-S2B5cJiosx_LpApd227w33CVvs', 'locked': False},
- {'key': 'AIzaSyCor0q5w37Dy6fGxloLlCT7KqyEFU3PWP8', 'locked': False},
- ]
- lock = asyncio.Lock()
- class VideoRequest(BaseModel):
- video_path: str
- prompt: str
- mark: str
- sample_data: str
- async def get_available_api_key():
- """获取一个未锁定的 API key,如果没有可用的则等待 5 秒后重试"""
- while True:
- async with lock:
- for key_data in api_keys:
- if not key_data['locked']:
- key_data['locked'] = True # 锁定该 key
- return key_data['key']
- print( "没有可用的 API key,等待 5 秒后重试..." )
- await asyncio.sleep( 5 )
- async def release_api_key(api_key):
- """释放已锁定的 API key,并将其放到列表末尾"""
- async with lock:
- for i, key_data in enumerate( api_keys ):
- if key_data['key'] == api_key:
- key_data['locked'] = False # 释放该 key
- # 将释放的 key 移动到列表末尾
- api_keys.append( api_keys.pop( i ) )
- break
- @app.post("/process_video/")
- async def process_video(request: VideoRequest):
- """处理视频请求"""
- global api_key_index # 引用全局索引
- video_path = request.video_path
- prompt = request.prompt
- mark = request.mark
- sample_data = request.sample_data
- api_key = await get_available_api_key()
- try:
- print("来一个请求,使用 API key:", api_key)
- result = await main(video_path, api_key, prompt, sample_data)
- return {
- "code": 0,
- "message": "视频处理成功",
- "result": result,
- "mark": str(mark)
- }
- except Exception as e:
- print(f"处理失败: {str(e)}")
- return {
- "code": 1,
- "message": f"处理失败: {e}",
- "result": f"处理失败: {e}",
- "mark": f"处理失败: {e}"
- }
- finally:
- await release_api_key( api_key )
- if __name__ == "__main__":
- uvicorn.run(app, host="0.0.0.0", port=8080)
|