@@ -1,120 +1,40 @@
-import logging
-import time
-import uuid
-from datetime import timedelta
-from decimal import Decimal
-from io import BytesIO
 from typing import Any
 
-import httpx
-from django.conf import settings
-from django.core.files import File
-
 from messages.models import Message
-from ml_model.exceptions import ModelVersionNotAvailable, RequestBlocked
-from ml_model.services.base import SimpleService
-from poller.models import Proxy
-
-logger = logging.getLogger(__name__)
+from ml_model.exceptions import ModelVersionNotAvailable
+from ml_model.services.seedream import Seedream
 
 
-class Stablediffusion(SimpleService):
+class Stablediffusion(Seedream):
     """
     Stablediffusion Service
     contains abstract method make, which makes a generation
     """
 
-    MODELS = ['sd3', 'sd3-turbo', 'sd3-medium']
-    MODELS_LINKS = {
-        'sd3': 'stable-diffusion-3.5-large',
-        'sd3-turbo': 'stable-diffusion-3.5-large-turbo',
-        'sd3-medium': 'stable-diffusion-3.5-medium',
+    PROXY_VERSION_MAPPING = {
+        'sd3': 'seedream-boosted',
+        'sd3-medium': 'seedream-4.5',
+        'sd3-turbo': 'seedream-4.5',
     }
 
-    def calculate_price(self, input_message: Message) -> Decimal:
-        if input_message.info.get('version') == 'sd3':
-            return Decimal('32.5')
-        elif input_message.info.get('version') == 'sd3-turbo':
-            return Decimal('20')
-        elif input_message.info.get('version') == 'sd3-medium':
-            return Decimal('17.5')
-
     @classmethod
-    def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None:
-        version = info.get('version')
-        if version == 'sd3':
-            return Decimal('32.5')
-        elif version == 'sd3-turbo':
-            return Decimal('20')
-        elif version == 'sd3-medium':
-            return Decimal('17.5')
-        return None
+    def _remap_info(cls, info: dict[str, Any]) -> dict[str, Any]:
+        new_info = info.copy()
+        version = info.get('version', 'sd3')
 
-    def save_results(
-        self,
-        input_prompt: str,
-        link: str,
-        t: timedelta,
-        save: bool = True,
-    ) -> list[Message]:
-        out: list[Message] = []
-        out.append(
-            Message(
-                content_object=self.store,
-                elapsed_time=t,
-                content=input_prompt,
-                file=File(
-                    BytesIO(httpx.get(link).content),
-                    f'{uuid.uuid4()}.png',
-                ),
-            )
-        )
-        if save:
-            return Message.objects.bulk_create(out)
-        return out
+        if version not in cls.PROXY_VERSION_MAPPING and version not in cls.PROXY_VERSION_MAPPING.values():
+            raise ModelVersionNotAvailable(version, cls.PROXY_VERSION_MAPPING)
+        new_info['version'] = cls.PROXY_VERSION_MAPPING.get(version, version)
 
-    def make(self, input_message: Message, save: bool = True) -> list[Message]:
-        start_time = time.time()
-        info = input_message.info.copy()
-        version = info.get('version')
-        if version is None or version not in self.MODELS:
-            raise ModelVersionNotAvailable(version, self.MODELS)
-        model_name = self.MODELS_LINKS[version]
-        translated_prompt = self.translate_prompt(input_message.content)
-        callback_data = {
-            'prompt': translated_prompt,
-            'aspect_ratio': input_message.info.get('aspect_ratio', '1:1'),
-            'output_quality': 100,
-            'output_format': 'png',
-        }
-        link = None
-        for proxy in Proxy.objects.all():
-            with httpx.Client(
-                headers={
-                    'Authorization': f'Bearer {settings.REPLICATE_API_KEY}',
-                    'Prefer': 'wait',
-                    'Content-Type': 'application/json',
-                },
-                timeout=600,
-                proxy=f'{proxy.protocol}://{proxy.address}',
-            ) as client:
-                result = client.post(
-                    f'https://api.replicate.com/v1/models/stability-ai/{model_name}/predictions',
-                    json={'input': callback_data},
-                ).json()
-                while result['status'] not in ('succeeded', 'failed', 'canceled'):
-                    result = client.get(result['urls']['get']).json()
-                if result['status'] in ('failed', 'canceled'):
-                    if 'E005' in result['logs']:
-                        raise RequestBlocked
-                    logger.error(result['logs'])
-                    raise Exception('No answer from Stable Diffusion, please retry later')
-                link = result['output']
+        return new_info
 
-            process_time = timedelta(seconds=(time.time() - start_time))
+    @classmethod
+    def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]):
+        try:
+            return super().predict_price(content, file_exists, cls._remap_info(info))
+        except ModelVersionNotAvailable:
+            return None
 
-            if not link:
-                raise Exception('No answer from Stable Diffusion, please retry later')
-            self.handle_invoice(input_message.content_object.model, input_message)
-            msgs = self.save_results(input_message.content, link, process_time, save)
-            return msgs
+    def make(self, input_message: Message, save: bool = True) -> list[Message]:
+        input_message.info = self._remap_info(input_message.info)
+        return super().make(input_message, save)
\ No newline at end of file