Explorar el Código

Merge pull request #1644 from PoseidonLi0514/master

Image generation supports custom endpoint
master
zhayujie GitHub hace 10 meses
padre
commit
ac6e41abc8
No se encontró ninguna clave conocida en la base de datos para esta firma ID de clave GPG: 4AEE18F83AFDEB23
Se han modificado 1 ficheros con 3 adiciones y 1 borrados
  1. +3
    -1
      bot/openai/open_ai_image.py

+ 3
- 1
bot/openai/open_ai_image.py Ver fichero

@@ -12,16 +12,18 @@ from config import conf
class OpenAIImage(object):
def __init__(self):
openai.api_key = conf().get("open_ai_api_key")
openai.api_base = conf().get("open_ai_api_base")
if conf().get("rate_limit_dalle"):
self.tb4dalle = TokenBucket(conf().get("rate_limit_dalle", 50))

def create_img(self, query, retry_count=0, api_key=None):
def create_img(self, query, retry_count=0, api_key=None, api_base=None):
try:
if conf().get("rate_limit_dalle") and not self.tb4dalle.get_token():
return False, "请求太快了,请休息一下再问我吧"
logger.info("[OPEN_AI] image_query={}".format(query))
response = openai.Image.create(
api_key=api_key,
api_base=api_base,
prompt=query, # 图片描述
n=1, # 每次生成图片的数量
model=conf().get("text_to_image") or "dall-e-2",


Cargando…
Cancelar
Guardar