Fix pricing for multiple outputs

GooeyAI · Feb 20, 2024 · 1b3ca5b · 1b3ca5b
1 parent f82d5f4
commit 1b3ca5b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/recipes/CompareText2Img.py b/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)
diff --git a/recipes/GoogleGPT.py b/recipes/GoogleGPT.py
@@ -279,3 +279,6 @@ def run_v2(
             max_tokens=request.max_tokens,
             avoid_repetition=request.avoid_repetition,
         )
+
+    def get_raw_price(self, state: dict) -> float:
+        return self.price * state.get("num_outputs", 1)
diff --git a/recipes/VideoBots.py b/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")