@@ -0,0 +1,19 @@ +# Generated by Django 5.0.11 on 2025-05-16 21:22 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('authentication', '0016_alter_customusermodel_is_confirmed'), + ('ml_model', '0050_remove_modeltag_color'), + ] + + operations = [ + migrations.AddField( + model_name='businessuserhost', + name='private_models', + field=models.ManyToManyField(blank=True, related_name='private_models_hosts', to='ml_model.neuronmodel', verbose_name='Private models'), + ), + ] @@ -11,6 +11,7 @@ from authentication.models.business_group import BusinessGroup 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): @@ -78,9 +79,14 @@ class BusinessUserHost(BaseModel): allowed_models = ArrayField( models.CharField(max_length=64, blank=True), default=list, + blank=True, verbose_name=_('Allowed models'), ) + private_models = models.ManyToManyField( + NeuronModel, blank=True, verbose_name=_('Private models'), related_name='private_models_hosts' + ) + 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_models',) @admin.action(description='Показать баланс всем сотрудникам выбранных компаний') def show_balance_for_all(self, request, qs: QuerySet[BusinessUserHost]): @@ -330,6 +330,7 @@ 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') 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 "Туризм" @@ -1,6 +1,6 @@ from uuid import UUID -from django.db.models import Prefetch +from django.db.models import Prefetch, Q from django.utils.translation import gettext_lazy as _ from authentication.models.choices import InvitationStatus @@ -48,6 +48,11 @@ class NeuronModelSelector: if self.user.is_anonymous: return NeuronModelsSerializer(models, many=True) user_type = UserSelector(self.user).check_account_type() + models = NeuronModel.objects.filter( + Q(private_models_hosts__isnull=True) + | Q(private_models_hosts__accounts__user=self.user) + | Q(private_models_hosts__user=self.user) + ) 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,83 @@ +import time +import httpx + +from backend import settings +from decimal import Decimal +from datetime import timedelta +from io import BytesIO +from typing import Any + +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('17.5'), + }, + } + + 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') + + def save_results( + self, + prompt: str, + images: list[dict[str, Any]], + 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': 'fb90c17a-d410-41e7-9961-dc7c687bc627', + **input_message.info, + } + ) + 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, **callback_data}, + ).json() + while True: + status = client.get(result['status_url']).json() + if status.get('status') == 'COMPLETED': + break + time.sleep(1/3) + process_time = timedelta(seconds=(time.time() - start_time)) + final_result = client.get(result['response_url']).json() + images = [img['url'] for img in final_result['images']] + 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