Fix pricing for more recipes

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/DocSummary.py

@@ -181,6 +181,9 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
             case _:
                 raise NotImplementedError(f"{chain_type} not implemented")
 
+    def get_raw_price(self, state: dict) -> float:
+        return self.price * state.get("num_outputs", 1)
+
 
 MAP_REDUCE_PROMPT = """
 {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/GoogleImageGen.py

@@ -238,3 +238,6 @@ def render_example(self, state: dict):
 
     def preview_description(self, state: dict) -> str:
         return "Enter a Google Image Search query + your Img2Img text prompt describing how to alter the result to create a unique, relevant ai generated images for any search query."
+
+    def get_raw_price(self, state: dict) -> float:
+        return super().get_raw_price(state) * 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/SmartGPT.py

@@ -204,6 +204,9 @@ def render_steps(self):
     def preview_description(self, state: dict) -> str:
         return "SmartGPT is a cutting-edge AI technology that can be used to generate natural language responses to any given input. We have combined the power of [CoT](https://arxiv.org/abs/2305.02897), [Reflexion](https://arxiv.org/abs/2303.11366) & [DERA](https://arxiv.org/abs/2303.17071) into one pipeline so that you can use ChatGPT to its full potential! Input your prompt + a reflection/research prompt + a resolver prompt to use SmartGPT for enhanced text generation, natural language and incredible question-answer results."
 
+    def get_raw_price(self, state: dict) -> float:
+        return self.price * state.get("num_outputs", 1)
+
 
 def answers_as_prompt(texts: list[str], sep="\n\n") -> str:
     return sep.join(

recipes/Text2Audio.py

@@ -140,6 +140,9 @@ def render_example(self, state: dict):
     def preview_description(self, state: dict) -> str:
         return "Generate AI Music with text instruction prompts. AudiLDM is capable of generating realistic audio samples by process any text input. Learn more [here](https://huggingface.co/cvssp/audioldm-m-full)."
 
+    def get_raw_price(self, state: dict) -> float:
+        return super().get_raw_price(state) * state.get("num_outputs", 1)
+
 
 def _render_output(state):
     selected_models = state.get("selected_models", [])