open_ai_image.py 1.7KB

1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
1 jaar geleden
12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. import time
  2. import openai
  3. import openai.error
  4. from common.log import logger
  5. from common.token_bucket import TokenBucket
  6. from config import conf
  7. # OPENAI提供的画图接口
  8. class OpenAIImage(object):
  9. def __init__(self):
  10. openai.api_key = conf().get("open_ai_api_key")
  11. if conf().get("rate_limit_dalle"):
  12. self.tb4dalle = TokenBucket(conf().get("rate_limit_dalle", 50))
  13. def create_img(self, query, retry_count=0, api_key=None):
  14. try:
  15. if conf().get("rate_limit_dalle") and not self.tb4dalle.get_token():
  16. return False, "请求太快了,请休息一下再问我吧"
  17. logger.info("[OPEN_AI] image_query={}".format(query))
  18. response = openai.Image.create(
  19. api_key=api_key,
  20. prompt=query, # 图片描述
  21. n=1, # 每次生成图片的数量
  22. model=conf().get("text_to_image") or "dall-e-2",
  23. # size=conf().get("image_create_size", "256x256"), # 图片大小,可选有 256x256, 512x512, 1024x1024
  24. )
  25. image_url = response["data"][0]["url"]
  26. logger.info("[OPEN_AI] image_url={}".format(image_url))
  27. return True, image_url
  28. except openai.error.RateLimitError as e:
  29. logger.warn(e)
  30. if retry_count < 1:
  31. time.sleep(5)
  32. logger.warn("[OPEN_AI] ImgCreate RateLimit exceed, 第{}次重试".format(retry_count + 1))
  33. return self.create_img(query, retry_count + 1)
  34. else:
  35. return False, "画图出现问题,请休息一下再问我吧"
  36. except Exception as e:
  37. logger.exception(e)
  38. return False, "画图出现问题,请休息一下再问我吧"