@@ -1,77 +1,123 @@ import base64 import logging +import math import time from datetime import timedelta from decimal import Decimal from io import BytesIO +import filetype import httpx +import tiktoken from django.core.files import File +from django.core.files.images import get_image_dimensions from backend import settings from messages.models import BaseStore, Message -from ml_model.exceptions import RequestBlocked, GenerationException +from ml_model.exceptions import ( + CorruptedFileError, + FileExtensionNotSupported, + GenerationException, + RequestBlocked, +) from ml_model.services.base import SimpleService +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy class Gptimage(SimpleService): - """ - GPTImage Service - contains abstract method make, which makes a generation - """ - TOKENS_COST = { - 'gpt-image-1': { + 'gpt-image-2': { 'input': Decimal('0.0025'), 'image': { - 'input': Decimal('0.005'), - 'output': Decimal('0.02'), - 'generation': { - 'low': { - '1024x1024': Decimal('5.5'), - '1024x1536': Decimal('8'), - '1536x1024': Decimal('8'), - }, - 'medium': { - '1024x1024': Decimal('21'), - '1024x1536': Decimal('31.5'), - '1536x1024': Decimal('31.5'), - }, - 'high': { - '1024x1024': Decimal('83.5'), - '1024x1536': Decimal('125'), - '1536x1024': Decimal('125'), - }, - }, + 'input': Decimal('0.004'), + 'output': Decimal('0.015'), }, - } + }, + } + QUALITY = { + 'Низкое': 'low', + 'Высокое': 'medium', + } + SIZE_TOKENS = { + 'low': { + '1024x1024': 196, + '1536x1024': 158, + '1024x1536': 158, + }, + 'medium': { + '1024x1024': 1756, + '1536x1024': 1372, + '1024x1536': 1372, + }, + } + MODERATION = { + 'Авто': 'auto', + 'Низкая': 'low', } def __init__(self, store: BaseStore) -> None: super().__init__(store) self.logger = logging.getLogger(self.__class__.__name__) + def calculate_price(self, text_tokens: int, image_tokens: int, output_tokens: int) -> Decimal: + price = ( + text_tokens * self.TOKENS_COST['gpt-image-2']['input'] + + image_tokens * self.TOKENS_COST['gpt-image-2']['image']['input'] + + output_tokens * self.TOKENS_COST['gpt-image-2']['image']['output'] + ) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, prompt: str, image: bytes, time: timedelta, save: bool = True) -> list[Message]: + messages: list[Message] = [] + image_bytes = base64.b64decode(image) + messages.append( + Message( + content_object=self.store, + elapsed_time=time, + content=prompt, + file=File(BytesIO(image_bytes), '.png'), + ) + ) + if save: + return Message.objects.bulk_create(messages) + return messages + def make(self, input_message: Message, save: bool = True) -> list[Message]: info = input_message.info.copy() - files = None + size = info.get('size', '1024x1024') + quality = self.QUALITY[info.get('quality', 'Низкое')] + w, h = 0, 0 + if file := input_message.file: + file_bytes = file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + if kind.extension.upper() not in (extensions := ['PNG', 'JPG', 'JPEG', 'WEBP']): + raise FileExtensionNotSupported(extensions) + file.seek(0) + w, h = get_image_dimensions(file) + if not w or not h: + raise CorruptedFileError + files = {'image': ('image.png', BytesIO(file_bytes), 'image/png')} + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.calculate_price( + *self.count_predict_tokens(input_message.content, w, h, size, quality) + ) + + (1 if not input_message.file else 4) + ): + raise InsufficientBalance(balance, cost) for proxy in Proxy.objects.all(): json_data = { 'prompt': input_message.content, - 'background': info.get('background', 'auto'), - 'model': 'gpt-image-1', + 'model': 'gpt-image-2', + 'n': 1, + 'quality': quality, + 'size': size, + 'moderation': self.MODERATION[info.get('moderation', 'Авто')], 'output_format': 'png', - 'quality': info.get('quality', 'auto'), - 'size': info.get('size', 'auto'), } - if input_message.file: - version = 'image-edit' - file_bytes = input_message.file.read() - files = { - 'image': ('image.png', BytesIO(file_bytes), 'image/png'), - } - else: - version = 'image-generate' with httpx.Client( base_url='https://api.openai.com/v1', proxy=f'{proxy.protocol}://{proxy.address}', @@ -79,7 +125,7 @@ class Gptimage(SimpleService): timeout=600, ) as client: start_time = time.time() - if version == 'image-generate': + if not input_message.file: data = client.post('images/generations', json=json_data).json() else: data = client.post('images/edits', data=json_data, files=files).json() @@ -99,38 +145,25 @@ class Gptimage(SimpleService): msgs = self.save_results(input_message.content, image, process_time, save) return msgs - def calculate_price( - self, - text_tokens: int, - image_tokens: int, - output_tokens: int, - *args, - **kwargs, - ) -> Decimal: - price = ( - text_tokens * self.TOKENS_COST['gpt-image-1']['input'] - + image_tokens * self.TOKENS_COST['gpt-image-1']['image']['input'] - + output_tokens * self.TOKENS_COST['gpt-image-1']['image']['output'] - ) - return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def count_predict_tokens( + cls, text_input: str, image_width: int, image_height: int, size: str, quality: str + ): + encoding = tiktoken.get_encoding('o200k_base') + text_input_tokens = len(encoding.encode(text_input)) + image_input_tokens = cls.count_image_tokens(image_width, image_height) + image_output_tokens = cls.SIZE_TOKENS[quality][size] + return text_input_tokens, image_input_tokens, image_output_tokens - def save_results( - self, - prompt: str, - image: bytes, - time: timedelta, - save: bool = True, - ) -> list[Message]: - messages: list[Message] = [] - image_bytes = base64.b64decode(image) - messages.append( - Message( - content_object=self.store, - elapsed_time=time, - content=prompt, - file=File(BytesIO(image_bytes), '.png'), - ) - ) - if save: - return Message.objects.bulk_create(messages) - return messages + @staticmethod + def count_image_tokens(w: int, h: int) -> int: + try: + p = math.ceil(w / 16) * math.ceil(h / 16) + r = max(w, h) / min(w, h) + s = 1521 * (1 - math.exp(-p / 1150.0)) + a = math.exp(-0.16 * (math.log(r) ** 2)) + k = 0.72 + 0.28 * (1 - math.exp(-p / 900.0)) + t = s * a * k + return int(round(max(16, min(1521, t)))) + except ZeroDivisionError: + return 0