analyze_video.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import uvicorn
  2. import asyncio
  3. from fastapi import FastAPI
  4. from pydantic import BaseModel
  5. from google_ai.generativeai_video import main
  6. import google.generativeai as genai
  7. app = FastAPI()
  8. api_keys = [
  9. {'key': 'AIzaSyB2kjF2-S2B5cJiosx_LpApd227w33CVvs', 'locked': False},
  10. {'key': 'AIzaSyCor0q5w37Dy6fGxloLlCT7KqyEFU3PWP8', 'locked': False},
  11. ]
  12. lock = asyncio.Lock()
  13. class VideoRequest(BaseModel):
  14. video_path: str
  15. prompt: str
  16. mark: str
  17. sample_data: str
  18. async def get_available_api_key():
  19. """获取一个未锁定的 API key,如果没有可用的则等待 5 秒后重试"""
  20. while True:
  21. async with lock:
  22. for key_data in api_keys:
  23. if not key_data['locked']:
  24. key_data['locked'] = True # 锁定该 key
  25. return key_data['key']
  26. print( "没有可用的 API key,等待 5 秒后重试..." )
  27. await asyncio.sleep( 5 )
  28. async def release_api_key(api_key):
  29. """释放已锁定的 API key,并将其放到列表末尾"""
  30. async with lock:
  31. for i, key_data in enumerate( api_keys ):
  32. if key_data['key'] == api_key:
  33. key_data['locked'] = False # 释放该 key
  34. # 将释放的 key 移动到列表末尾
  35. api_keys.append( api_keys.pop( i ) )
  36. break
  37. @app.post("/process_video/")
  38. async def process_video(request: VideoRequest):
  39. """处理视频请求"""
  40. global api_key_index # 引用全局索引
  41. video_path = request.video_path
  42. prompt = request.prompt
  43. mark = request.mark
  44. sample_data = request.sample_data
  45. api_key = await get_available_api_key()
  46. try:
  47. print("来一个请求,使用 API key:", api_key)
  48. result = await main(video_path, api_key, prompt, sample_data)
  49. return {
  50. "code": 0,
  51. "message": "视频处理成功",
  52. "result": result,
  53. "mark": str(mark)
  54. }
  55. except Exception as e:
  56. print(f"处理失败: {str(e)}")
  57. return {
  58. "code": 1,
  59. "message": f"处理失败: {e}",
  60. "result": f"处理失败: {e}",
  61. "mark": f"处理失败: {e}"
  62. }
  63. finally:
  64. await release_api_key( api_key )
  65. if __name__ == "__main__":
  66. uvicorn.run(app, host="0.0.0.0", port=8080)