| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364 |
- from typing import Dict
- from openai import OpenAI
- from common.config.tokenizer_manage_config import TokenizerManage
- from models_provider.base_model_provider import MaxKBBaseModel
- from models_provider.impl.base_tti import BaseTextToImage
- def custom_get_token_ids(text: str):
- tokenizer = TokenizerManage.get_tokenizer()
- return tokenizer.encode(text)
- class OpenAITextToImage(MaxKBBaseModel, BaseTextToImage):
- api_base: str
- api_key: str
- model: str
- params: dict
- def __init__(self, **kwargs):
- super().__init__(**kwargs)
- self.api_key = kwargs.get('api_key')
- self.api_base = kwargs.get('api_base')
- self.model = kwargs.get('model')
- self.params = kwargs.get('params')
- @staticmethod
- def is_cache_model():
- return False
- @staticmethod
- def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
- optional_params = {'params': {'size': '1024x1024', 'quality': 'standard', 'n': 1}}
- for key, value in model_kwargs.items():
- if key not in ['model_id', 'use_local', 'streaming']:
- optional_params['params'][key] = value
- return OpenAITextToImage(
- model=model_name,
- api_base=model_credential.get('api_base'),
- api_key=model_credential.get('api_key'),
- **optional_params,
- )
- def check_auth(self):
- chat = OpenAI(api_key=self.api_key, base_url=self.api_base)
- response_list = chat.models.with_raw_response.list()
- # self.generate_image('生成一个小猫图片')
- def generate_image(self, prompt: str, negative_prompt: str = None):
- chat = OpenAI(api_key=self.api_key, base_url=self.api_base)
- res = chat.images.generate(model=self.model, prompt=prompt, **self.params)
- file_urls = []
- try:
- for content in res.data:
- if content.url:
- file_urls.append(content.url)
- elif content.b64_json:
- file_urls.append(content.b64_json)
- return file_urls
- except Exception as e:
- raise f"OpenAI generate image error: {e}"
|
