@@ -1,4 +1,3 @@
-import base64
 import time
 from datetime import timedelta
 from decimal import Decimal
@@ -11,43 +10,37 @@ from replicate.exceptions import ModelError
 
 from messages.models import Message
 from ml_model.exceptions import RequestBlocked, GenerationException
-from ml_model.models import (
-    NeuronModel,
-)
 from ml_model.services.base import SimpleService
 from ml_model.tasks import replicate_run
 
 
 class Leonardo(SimpleService):
-    """
-    Flux Service
-    contains abstract method make, which makes a generation
-    """
-
     TOKENS_COST = {
         'lucid-origin': {
-            'input_imgs': Decimal('1.5'),
-        }, # 1k images
+            'input_units': Decimal('450'),
+        }
     }
 
     _CALLBACK_BASE = 'leonardoai/'
 
-    def calculate_price(self, input_message: Message, version: str) -> Decimal:
-        price_map = self.TOKENS_COST[version]
-        price = price_map['input_imgs'] / 1_000
-        if num_images := input_message.info.get('num_images'):
-            price = price * num_images
+    def calculate_price(self, version: str, num_images: int, generation_mode: str) -> Decimal:
+        image_prices = {'standard': 11, 'ultra': 51}
+        price = (
+            Decimal(f'{self.TOKENS_COST[version]["input_units"] / 1000 * image_prices[generation_mode]}')
+            * num_images
+        )
         return price.quantize(Decimal('0.1'), rounding='ROUND_UP')
 
     @classmethod
     def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None:
-        price = cls.TOKENS_COST['lucid-origin']['input_imgs'] / 1_000 * info['num_images']
+        version = info.get('version', 'lucid-origin')
+        num_images = info.get('num_images', 1)
+        generation_mode = info.get('generation_mode', 'standard')
+        image_prices = {'standard': 11, 'ultra': 51}
+        input_units = cls.TOKENS_COST[version]['input_units']
+        price = (input_units / 1000 * image_prices[generation_mode]) * num_images
         return price.quantize(Decimal('0.1'), rounding='ROUND_UP')
 
-    @property
-    def neuron_model(self):
-        return NeuronModel.objects.get(title='Flux')
-
     def save_results(
         self,
         prompt: str,
@@ -71,7 +64,9 @@ class Leonardo(SimpleService):
 
     def make(self, input_message: Message, save: bool = True) -> list[Message]:
         start_time = time.time()
-        version = input_message.info.get('version')
+        version = input_message.info.get('version', 'lucid-origin')
+        generation_mode = input_message.info.get('generation_mode', 'standard')
+        num_images = input_message.info.get('num_images', 1)
         callback_data = dict(
             {
                 'prompt': self.translate_prompt(input_message.content),
@@ -89,6 +84,11 @@ class Leonardo(SimpleService):
             raise GenerationException from exc
         images = runner if isinstance(runner, list) else [runner]
         process_time = timedelta(seconds=(time.time() - start_time))
-        self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version)
+        self.handle_invoice(
+            input_message.content_object.model,
+            version=version,
+            num_images=num_images,
+            generation_mode=generation_mode,
+        )
         msgs = self.save_results(input_message.content, images, process_time, save)
         return msgs