@@ -44,8 +44,10 @@ class Prunaai(SimpleService):
 
     def make(self, input_message: Message, save: bool = True) -> list[Message]:
         version = input_message.info.get('version', 'p-image')
+        aspect_ratio = input_message.info.pop('aspect_ratio', 'custom')
         if input_message.file:
             version = 'p-image-edit'
+            aspect_ratio = 'match_input_image'
         if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST[
             version
         ]:
@@ -53,6 +55,7 @@ class Prunaai(SimpleService):
         callback_data = dict(
             {
                 'prompt': self.translate_prompt(input_message.content),
+                'aspect_ratio': aspect_ratio,
                 **input_message.info,
             }
         )