@@ -12,6 +12,8 @@ from authentication.models.choices import BusinessSector, UsageIntensity from authentication.models.whitelist import CompanyIPWhitelist from core.models import BaseModel +from ml_model.models import NeuronModel + class BusinessUserHost(BaseModel): user = models.OneToOneField( @@ -81,6 +83,8 @@ class BusinessUserHost(BaseModel): verbose_name=_('Allowed models'), ) + private = models.ManyToManyField(NeuronModel, blank=True, verbose_name=_('Private models'), related_name='private_models') + is_log_history_enabled = models.BooleanField(default=False, verbose_name=_('Log history enabled')) @property @@ -363,6 +363,7 @@ class BusinessUserHostAdmin(admin.ModelAdmin): verbose_name_plural = 'Бизнес-аккаунты-хосты' raw_id_fields = ['user', 'affiliated_by'] actions = ['show_balance_for_all', 'hide_balance_for_all'] + filter_horizontal = ('private',) @admin.action(description='Показать баланс всем сотрудникам выбранных компаний') def show_balance_for_all(self, request, qs: QuerySet[BusinessUserHost]): @@ -330,6 +330,8 @@ GOOGLE_API_KEY = env.str('GOOGLE_API_KEY', default='defaultapikey') SERPER_API_KEY = env.str('SERPER_API_KEY', 'defaultapikey') FLUX_API_KEY = env.str('FLUX_API_KEY', 'defaultapikey') OPENROUTER_API_KEY = env.str('OPENROUTER_API_KEY', 'defaultapikey') +FAL_API_KEY = env.str('FAL_API_KEY', 'defaultapikey') +FAL_API_MODEL = env.str('FAL_API_MODEL', 'defaultapikey') OPENAI_PROXY_HOST = env.str('OPENAI_PROXY_HOST', 'neuron-proxy:8080') UPSCALE_MULTIPLIER_HOST = env.str('UPSCALE_MULTIPLIER_HOST', 'packet:8080') @@ -250,6 +250,10 @@ msgstr "IT" msgid "Finance" msgstr "Финансы" +#: authentication/models/business_host.py:86 +msgid "Private models" +msgstr "Приватные модели" + #: authentication/models/choices.py:9 msgid "Tourism" msgstr "Туризм" @@ -3,6 +3,7 @@ from uuid import UUID from django.db.models import Prefetch from django.utils.translation import gettext_lazy as _ +from authentication.models import BusinessUserHost, BusinessAccount from authentication.models.choices import InvitationStatus from authentication.models.user import CustomUserModel from authentication.selectors.user_selector import UserSelector @@ -48,6 +49,16 @@ class NeuronModelSelector: if self.user.is_anonymous: return NeuronModelsSerializer(models, many=True) user_type = UserSelector(self.user).check_account_type() + for host in BusinessUserHost.objects.prefetch_related('private'): + for private_model in host.private.all(): + if user_type == 'regular': + models = models.exclude(title=private_model) + else: + user_private_models = set() + for name_host in BusinessUserHost.objects.filter(accounts__user=self.user): + user_private_models.update(name_host.private.all()) + if private_model not in user_private_models: + models = models.exclude(title=private_model) if ( user_type == 'business_account' and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED @@ -8,6 +8,7 @@ from ml_model.services.djourney import Djourney from ml_model.services.epicphotogasm import Epicphotogasm from ml_model.services.flux import Flux from ml_model.services.fluxproultra import Fluxproultra +from ml_model.services.fluxlorafast import Fluxlorafast from ml_model.services.gemini import Gemini from ml_model.services.granite import Granite from ml_model.services.grok import Grok @@ -0,0 +1,96 @@ +import base64 +import filetype +import time +import httpx + +from backend import settings +from decimal import Decimal +from datetime import timedelta +from io import BytesIO + +from django.core.files import File + +from messages.models import Message + +from ml_model.models import ModelInput +from ml_model.services.base import SimpleService + + +class Fluxlorafast(SimpleService): + TOKENS_COST = { + 'flux-lora': { + 'input_imgs': Decimal('2'), # узнать прайс надо + }, + } + + def calculate_price(self, input_message: Message, version: str) -> Decimal: + price_map = self.TOKENS_COST[version] + price = price_map['input_imgs'] + if image_count := input_message.info.get('num_outputs'): + price = price * image_count + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + inputs = [ + ModelInput(type=ModelInput.TypeChoices.TEXT, required=True), + ModelInput(type=ModelInput.TypeChoices.IMAGE), + ] + + def save_results( + self, + prompt: str, + images: list, + time: timedelta, + save: bool = True, + ) -> list[Message]: + messages: list[Message] = [] + for image in images: + messages.append( + Message( + content_object=self.store, + elapsed_time=time, + content=prompt, + file=File(BytesIO(httpx.get(image).content), '.png'), + ) + ) + if save: + return Message.objects.bulk_create(messages) + return messages + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + version = input_message.info.get('version') + translated_prompt = self.translate_prompt(input_message.content) + callback_data = dict( + { + 'prompt': translated_prompt, + 'model_version': settings.FAL_API_MODEL, + **input_message.info, + } + ) + if input_message.file: + kind = filetype.guess(input_message.file.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + input_message.file.seek(0) + image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' + input_message.file.close() + callback_data.update({'image': image}) + client = httpx.Client( + base_url="https://queue.fal.run", + headers={"Authorization": f"Key {settings.FAL_API_KEY}"}, + timeout=600, + ) + result = client.post( + f'fal-ai/{version}', + json={'prompt': input_message.content}, + ).json() + while True: + status = client.get(result['status_url']).json() + if status.get('status') == 'COMPLETED': + break + time.sleep(2) + final_result = client.get(result['response_url']).json() + images = [img['url'] for img in final_result['images']] + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version) + msgs = self.save_results(input_message.content, images, process_time, save) + return msgs