@@ -29,12 +29,12 @@ class Stablediffusion(SimpleService):
         'sd3-medium': 'stable-diffusion-3.5-medium'
     }
 
-    def calculate_price(self, input_message: Message) -> Decimal:
-        if input_message.info.get('version') == 'sd3':
+    def calculate_price(self, version: str) -> Decimal:
+        if version == 'sd3':
             return Decimal('32.5')
-        elif input_message.info.get('version') == 'sd3-turbo':
+        elif version == 'sd3-turbo':
             return Decimal('20')
-        elif input_message.info.get('version') == 'sd3-medium':
+        elif version == 'sd3-medium':
             return Decimal('17.5')
 
     def save_results(
@@ -63,7 +63,8 @@ class Stablediffusion(SimpleService):
     def make(self, input_message: Message, save: bool = True) -> list[Message]:
         start_time = time.time()
         info = input_message.info.copy()
-        model_name = self.MODELS_LINKS[info.get('version', 'sd3')]
+        version = info.get('version', 'sd3')
+        model_name = self.MODELS_LINKS[version]
         translated_prompt = self.translate_prompt(input_message.content)
         callback_data = {
             'prompt': translated_prompt,
@@ -92,8 +93,7 @@ class Stablediffusion(SimpleService):
                     logger.error(result.json()['logs'])
                     raise Exception('No answer from Stable Diffusion, please retry later')
                 link = result.json()['output'][0]
-
             process_time = timedelta(seconds=(time.time() - start_time))
-            self.handle_invoice(input_message.content_object.model, input_message)
+            self.handle_invoice(input_message.content_object.model, version)
             msgs = self.save_results(input_message.content, link, process_time, save)
             return msgs