| 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() |
| 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}", |
| } |
