@@ -1,25 +1,32 @@
-import base64
 import time
 from datetime import timedelta
 from decimal import Decimal
 from io import BytesIO
 from typing import Any
 
-import filetype
 import requests
 from django.core.files import File
 
 from messages.models import Message
 from ml_model.adapters.bytedance_model_ark import BytedanceContentType
+from ml_model.exceptions import InvalidParameterError
 from ml_model.services.base import SimpleService
-from ml_model.tasks import bytedance_model_ark_run, replicate_run
+from ml_model.tasks import bytedance_model_ark_run
+
+# import base64
+# import filetype
+# from ml_model.tasks import replicate_run
 
 
 class Reve(SimpleService):
     # Reve временно не работает на репликейте. Временно используем сидрим
     TEMPORARY_PROVIDER_MODEL = 'seedream-5-0-260128'
 
-    PRICE = Decimal('25')
+    PRICE = {
+        '2K': Decimal('25'),
+        '3K': Decimal('50'),
+        '4K': Decimal('100'),
+    }
     # PRICE = {
     #     'create': Decimal('12.5'),
     #     'edit-fast': Decimal('5'),
@@ -27,10 +34,12 @@ class Reve(SimpleService):
 
     @classmethod
     def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None:
-        return cls.PRICE.quantize(Decimal('0.1'), rounding='ROUND_UP')
+        price = cls.PRICE.get(info.get('size', '2K'))
+
+        return price.quantize(Decimal('0.1'), rounding='ROUND_UP')
 
-    def calculate_price(self) -> Decimal:
-        return self.PRICE.quantize(Decimal('0.1'), rounding='ROUND_UP')
+    def calculate_price(self, size: str) -> Decimal:
+        return self.PRICE[size].quantize(Decimal('0.1'), rounding='ROUND_UP')
 
     # @classmethod
     # def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None:
@@ -62,10 +71,14 @@ class Reve(SimpleService):
 
     def make(self, input_message: Message, save: bool = True) -> list[Message]:
         start_time = time.time()
+        size = input_message.info.get('size', '2K')
+        if size not in self.PRICE:
+            raise InvalidParameterError(f'Unsupported size: {size}')
+
         callback_data = {
             'prompt': self.translate_prompt(input_message.content),
             **input_message.info,
-            'size': '2K',
+            'size': size,
             'watermark': False,
         }
         if image := input_message.file:
@@ -76,7 +89,7 @@ class Reve(SimpleService):
             content_type=BytedanceContentType.IMAGE,
         )[0]
         process_time = timedelta(seconds=(time.time() - start_time))
-        self.handle_invoice(input_message.content_object.model)
+        self.handle_invoice(input_message.content_object.model, size=size)
         msgs = self.save_results(input_message.content, image, process_time, save)
         return msgs