Fix pricing for num_outputs>1

daras_ai_v2/base.py

@@ -1838,7 +1838,7 @@ def get_price_roundoff(self, state: dict) -> int:
         return max(1, math.ceil(self.get_raw_price(state)))
 
     def get_raw_price(self, state: dict) -> float:
-        return self.price
+        return self.price * state.get("num_outputs", 1)
 
     def get_example_response_body(
         self,

recipes/CompareText2Img.py

@@ -264,4 +264,4 @@ def get_raw_price(self, state: dict) -> int:
                     total += 15
                 case _:
                     total += 2
-        return total
+        return total * state.get("num_outputs", 1)

recipes/DocSearch.py

@@ -205,9 +205,10 @@ def run_v2(
     def get_raw_price(self, state: dict) -> float:
         name = state.get("selected_model")
         try:
-            return llm_price[LargeLanguageModels[name]] * 2
+            unit_price = llm_price[LargeLanguageModels[name]] * 2
         except KeyError:
-            return 10
+            unit_price = 10
+        return unit_price * state.get("num_outputs", 1)
 
 
 def render_documents(state, label="**Documents**", *, key="documents"):

recipes/FaceInpainting.py

@@ -336,6 +336,8 @@ def get_raw_price(self, state: dict) -> int:
         selected_model = state.get("selected_model")
         match selected_model:
             case InpaintingModels.dall_e.name:
-                return 20
+                unit_price = 20
             case _:
-                return 5
+                unit_price = 5
+
+        return unit_price * state.get("num_outputs", 1)

recipes/Img2Img.py

@@ -202,6 +202,8 @@ def get_raw_price(self, state: dict) -> int:
         selected_model = state.get("selected_model")
         match selected_model:
             case Img2ImgModels.dall_e.name:
-                return 20
+                unit_price = 20
             case _:
-                return 5
+                unit_price = 5
+
+        return unit_price * state.get("num_outputs", 1)

recipes/ObjectInpainting.py

@@ -314,6 +314,8 @@ def get_raw_price(self, state: dict) -> int:
         selected_model = state.get("selected_model")
         match selected_model:
             case InpaintingModels.dall_e.name:
-                return 20
+                unit_price = 20
             case _:
-                return 5
+                unit_price = 5
+
+        return unit_price * state.get("num_outputs", 1)

recipes/VideoBots.py

@@ -601,11 +601,13 @@ def get_raw_price(self, state: dict):
                     "raw_tts_text", state.get("raw_output_text", [])
                 )
                 tts_state = {"text_prompt": "".join(output_text_list)}
-                return super().get_raw_price(state) + TextToSpeechPage().get_raw_price(
+                total = super().get_raw_price(state) + TextToSpeechPage().get_raw_price(
                     tts_state
                 )
             case _:
-                return super().get_raw_price(state)
+                total = super().get_raw_price(state)
+
+        return total * state.get("num_outputs", 1)
 
     def additional_notes(self):
         tts_provider = st.session_state.get("tts_provider")