File size: 2,407 Bytes
1ea2ba0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import logging
from .base_model import BaseLLMModel
from .. import shared
import requests
from ..presets import *
from ..config import retrieve_proxy, sensitive_id
class OpenAI_DALLE3_Client(BaseLLMModel):
def __init__(self, model_name, api_key, user_name="") -> None:
super().__init__(model_name=model_name, user=user_name)
self.api_key = api_key
self._refresh_header()
def _get_dalle3_prompt(self):
prompt = self.history[-1]["content"]
if prompt.endswith("--raw"):
prompt = "I NEED to test how the tool works with extremely simple prompts. DO NOT add any detail, just use it AS-IS:" + prompt
return prompt
def get_answer_at_once(self, stream=False):
prompt = self._get_dalle3_prompt()
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}"
}
payload = {
"model": "dall-e-3",
"prompt": prompt,
"n": 1,
"size": "1024x1024",
"quality": "standard",
}
if stream:
timeout = TIMEOUT_STREAMING
else:
timeout = TIMEOUT_ALL
if shared.state.images_completion_url != IMAGES_COMPLETION_URL:
logging.debug(f"使用自定义API URL: {shared.state.images_completion_url}")
with retrieve_proxy():
try:
response = requests.post(
shared.state.images_completion_url,
headers=headers,
json=payload,
stream=stream,
timeout=timeout,
)
response.raise_for_status() # 根据HTTP状态码引发异常
response_data = response.json()
image_url = response_data['data'][0]['url']
img_tag = f'<!-- S O PREFIX --><a data-fancybox="gallery" target="_blank" href="{image_url}"><img src="{image_url}" /></a><!-- E O PREFIX -->'
revised_prompt = response_data['data'][0].get('revised_prompt', '')
return img_tag + revised_prompt, 0
except requests.exceptions.RequestException as e:
return str(e), 0
def _refresh_header(self):
self.headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {sensitive_id}",
} |