@@ -9,6 +9,7 @@ from ml_model.services.epicphotogasm import Epicphotogasm
 from ml_model.services.flux import Flux
 from ml_model.services.fluxlorafast import Fluxlorafast
 from ml_model.services.fluxproultra import Fluxproultra
+from ml_model.services.flux_2 import Flux_2
 from ml_model.services.gemini import Gemini
 from ml_model.services.geminiimage import Geminiimage
 from ml_model.services.gptimage import Gptimage
@@ -0,0 +1,92 @@
+import base64
+import math
+import time
+from datetime import timedelta
+from decimal import Decimal
+from io import BytesIO
+
+import filetype
+import requests
+from django.core.files import File
+from django.core.files.images import get_image_dimensions
+
+from messages.models import Message
+from ml_model.services.base import SimpleService
+from ml_model.tasks import replicate_run
+
+
+class Flux_2(SimpleService):
+    TOKENS_COST = {
+        'flux-2-pro': {
+            'run': Decimal('4.5'),
+            'input_mp': Decimal('4.5'),
+            'output_mp': Decimal('4.5'),
+        },
+        'flux-2-dev': {
+            'input_mp': Decimal('3.6'),
+            'output_mp': Decimal('3.6'),
+        },
+        'flux-2-flex': {
+            'input_mp': Decimal('18'),
+            'output_mp': Decimal('18'),
+        },
+    }
+
+    def calculate_price(self, version: str, input_mp: int, output_mp: int) -> Decimal:
+        version_price = self.TOKENS_COST[version]
+        price = (
+            version_price.get('run', Decimal('0')) +
+            version_price['input_mp'] * input_mp +
+            version_price['output_mp'] * output_mp
+        )
+        return price.quantize(Decimal('0.1'), rounding='ROUND_UP')
+
+    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(requests.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', 'flux-2-pro')
+        width = input_message.info.pop('width', 1024)
+        height = input_message.info.pop('height', 1024)
+        output_mp = math.ceil((width*height) / 1_000_000)
+        callback_data = {
+            'prompt': self.translate_prompt(input_message.content),
+            'aspect_ratio': 'custom',
+            'width': width,
+            'height': height,
+            **input_message.info,
+        }
+        input_mp = 0
+        if input_message.file:
+            file_width, file_height = get_image_dimensions(input_message.file)
+            input_mp = math.ceil((file_width*file_height) / 1_000_000)
+            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({'input_images': [image]})
+        images = [replicate_run(f'black-forest-labs/{version}', callback_data)]
+        process_time = timedelta(seconds=(time.time() - start_time))
+        self.handle_invoice(input_message.content_object.model, version=version, input_mp=input_mp, output_mp=output_mp)
+        msgs = self.save_results(input_message.content, images, process_time, save)
+        return msgs
@@ -3,6 +3,7 @@ import time
 from datetime import timedelta
 from decimal import Decimal
 from io import BytesIO
+from typing import Optional
 
 import filetype
 import requests
@@ -14,10 +15,19 @@ from ml_model.tasks import replicate_run
 
 
 class Nanobanana(SimpleService):
-    TOKENS_COST = Decimal('19.5')
+    TOKENS_COST = {
+        'nano-banana': Decimal('19.5'),
+        'nano-banana-pro': {
+            '1K': Decimal('45'),
+            '2K': Decimal('45'),
+            '4K': Decimal('90'),
+        }
+    }
 
-    def calculate_price(self) -> Decimal:
-        return self.TOKENS_COST
+    def calculate_price(self, version: str, resolution: Optional[str]) -> Decimal:
+        if resolution:
+            return self.TOKENS_COST[version][resolution]
+        return self.TOKENS_COST[version]
 
     def save_results(
         self,
@@ -38,7 +48,8 @@ class Nanobanana(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', 'nano-banana')
+        resolution = input_message.info.get('resolution', '2K') if version == 'nano-banana-pro' else None
         callback_data = dict(
             {
                 'prompt': self.translate_prompt(input_message.content),
@@ -52,8 +63,8 @@ class Nanobanana(SimpleService):
             image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}'
             input_message.file.close()
             callback_data.update({'image_input': [image]})
-        image = replicate_run('google/nano-banana', callback_data)
+        image = replicate_run(f'google/{version}', callback_data)
         process_time = timedelta(seconds=(time.time() - start_time))
-        self.handle_invoice(input_message.content_object.model)
+        self.handle_invoice(input_message.content_object.model, version=version, resolution=resolution)
         msgs = self.save_results(input_message.content, image, process_time, save)
         return msgs