@@ -1,6 +1,6 @@ from ml_model.runners.dummy import DummyImageRunner, DummyTextRunner from ml_model.runners.falai import FalAIRunner -from ml_model.runners.openai import OpenAIGPTRunner, OpenAIResponseRunner +from ml_model.runners.openai import OpenAIGPTRunner, OpenAIResponseRunner, GPTImageRunner from ml_model.runners.openrouter import OpenrouterRunner from ml_model.runners.replicate import ( ReplicateAudioRunner, @@ -12,6 +12,7 @@ from ml_model.runners.replicate import ( __all__ = [ 'OpenAIGPTRunner', 'OpenAIResponseRunner', + 'GPTImageRunner', 'OpenrouterRunner', 'ReplicateTextRunner', 'ReplicateAudioRunner', @@ -1,7 +1,6 @@ import base64 import json import logging -import re import uuid from abc import ABC, abstractmethod from io import BytesIO, StringIO @@ -77,9 +76,9 @@ class OpenAICompatibleRunner(BaseRunner, ABC): ) elif file and isinstance(file, StringIO): file_content = file.getvalue() - if len(file_data := file.getvalue()) > 20_000: + if len(file_content) > 20_000: chunks = TextSplitterTool().split_text( - text=file_data, separators=["\n\n", "\n", ".", " ", ""] + text=file_content, separators=["\n\n", "\n", ".", " ", ""] ) for proxy in proxies: try: @@ -289,3 +288,62 @@ class OpenAIResponseRunner(OpenAIGPTRunner): except httpx.TimeoutException as exc: logger.exception(exc) raise Exception('Timeout happened') + + +class GPTImageRunner(OpenAIGPTRunner): + @classmethod + def generate(cls, content=None, file=None, parameters={}, history=[], scrape_results=[]): + payload = { + 'prompt': content, + 'stream': True, + **parameters + } + + if file and isinstance(file, BytesIO): + files = [ + ('image[]', ('image.png', file, 'image/png')) + ] + + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {cls.AUTHORIZATION_TOKEN}', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + if file and isinstance(file, BytesIO): + streaming = client.stream('POST', 'images/edits', data=payload, files=files) + else: + streaming = client.stream('POST', 'images/generations', json=payload) + try: + with streaming as stream: + stream_content = stream.iter_lines() + if stream.status_code >= 400: + raw = ''.join([chunk for chunk in stream_content]) + try: + errors: dict[Literal['error'], dict[Literal['message'] | str, Any]] = ( + json.loads(raw).get('error', {}) + ) + except json.decoder.JSONDecodeError: + logger.error(raw) + errors = raw + cls.map_errors(errors) + for chunk in stream_content: + if chunk.strip() in cls.SKIP_TOKENS: + continue + if chunk.strip() in cls.END_TOKENS: + break + try: + dict_ = json.loads(chunk[5:].strip()) + if dict_.get("type") in ("image_edit.completed", "image_generation.completed"): + if b64_json := dict_.get("b64_json", ""): + yield b64_json + except json.decoder.JSONDecodeError: + continue + except httpx.ConnectError: + continue + except httpx.TimeoutException as exc: + logger.exception(exc) + raise Exception('Timeout happened') @@ -26,7 +26,8 @@ from messages.models import Message from ml_model.exceptions import ( InferenceDisabled, PaymentRuleNotImplemented, - ScraperDoesNotExists, UnknownFileException, + ScraperDoesNotExists, + UnknownFileException, ) from ml_model.models import ( Deployment, @@ -265,7 +266,7 @@ class InferenceService: ): if isinstance(file, StringIO) and len(file_data := file.getvalue()) >= 20_000: calculated_price += TokenizerTool.token_count( - file_data, biggest_coefficient=0.15 + file_data ) * payment_rule.cost elif ( content @@ -360,7 +361,6 @@ class InferenceService: process_time = timedelta(seconds=end - start) output_content = output_content.split('base64,')[-1] # cutoff b64-prefix if exists - output_slot.elapsed_time = process_time match inference.deployment.output_type: case Deployment.OutputTypeChoices.TEXT: @@ -379,7 +379,6 @@ class InferenceService: and payment_rule.interaction_type == PaymentRule.InteractionTypeChoices.OUTPUT ): calculated_price += TokenizerTool.token_count(output_content) * payment_rule.cost - for payment_bias in inference.payment_biases: if payment_bias.type == PaymentBias.TypeChoices.ADDITION: calculated_price += payment_bias.coefficient @@ -8,7 +8,7 @@ from typing import List from redis.commands.search.document import Document from redis.commands.search.query import Query -from .tasks import drop_redis_vectors +from ml_model.tools.tasks import drop_redis_vectors class EmbeddingTool: @@ -1,4 +1,4 @@ -import unicodedata +import tiktoken class TokenizerTool: @@ -8,35 +8,8 @@ class TokenizerTool: def token_count( self, text: str, - subword_step: int = 3, - smaller_coefficient: float = 0.15, - biggest_coefficient: float = 0.05, - token_bias: int = 0 ) -> int: - if not text: - return 0 - text = unicodedata.normalize("NFC", text) - count = 0 - i = 0 - while i < len(text): - c = text[i] - if c.isspace(): - count += 1 - i += 1 - elif c in '.,!?;:()[]{}"\'«»—–-0123456789': - count += 1 - i += 1 - else: - utf_bytes = c.encode('utf-8') - j = i + 1 - while j < len(text): - next_char = text[j] - if next_char.isspace() or next_char in '.,!?;:()[]{}"\'«»—–-0123456789': - break - utf_bytes += next_char.encode('utf-8') - j += 1 - count += max(1, len(utf_bytes) // subword_step) - i = j - if count < 100: - return int(round(count + count * smaller_coefficient) + token_bias) - return int(round(count / 2 + count * biggest_coefficient) + token_bias) \ No newline at end of file + encoding = tiktoken.get_encoding('o200k_base') + num_tokens = len(encoding.encode(text)) + return int(round(num_tokens + num_tokens * 0.20)) +