More Storybot updates (#1116)

* initial changes for gemini storybot

* storybot updates for gemini

* more storybot updates

* interim interruptible commit

* cleanup

* cleanup

* cleanup

* first draft

* wip

* more storybot fixes

* more storybot updates WIP

* committing before changing the image prompting strategy

* wip

* prompt updating

* cleanup

* cleanup

* cleanup

* readme cleanup

* fixup

examples/storytelling-chatbot/src/bot.py

@@ -51,8 +51,8 @@ async def main(room_url, token=None):
             DailyParams(
                 audio_out_enabled=True,
                 camera_out_enabled=True,
-                camera_out_width=768,
-                camera_out_height=768,
+                camera_out_width=1024,
+                camera_out_height=1024,
                 transcription_enabled=True,
                 vad_analyzer=SileroVADAnalyzer(),
                 vad_enabled=True,
@@ -69,16 +69,7 @@ async def main(room_url, token=None):
             api_key=os.getenv("ELEVENLABS_API_KEY"), voice_id=os.getenv("ELEVENLABS_VOICE_ID")
         )
 
-        fal_service_params = FalImageGenService.InputParams(
-            image_size={"width": 768, "height": 768}
-        )
-
-        fal_service = FalImageGenService(
-            aiohttp_session=session,
-            model="fal-ai/stable-diffusion-v35-medium",
-            params=fal_service_params,
-            key=os.getenv("FAL_KEY"),
-        )
+        image_gen = GoogleImageGenService(api_key=os.getenv("GOOGLE_API_KEY"))
 
         # --------------- Setup ----------------- #
 
@@ -92,14 +83,13 @@ async def main(room_url, token=None):
         # -------------- Processors ------------- #
 
         story_processor = StoryProcessor(message_history, story_pages)
-        image_processor = StoryImageProcessor(fal_service)
+        image_processor = StoryImageProcessor(image_gen)
 
         # -------------- Story Loop ------------- #
 
         runner = PipelineRunner()
 
         logger.debug("Waiting for participant...")
-
         main_pipeline = Pipeline(
             [
                 transport.input(),
@@ -119,7 +109,6 @@ async def main(room_url, token=None):
                 allow_interruptions=True,
                 enable_metrics=True,
                 enable_usage_metrics=True,
-                report_only_initial_ttfb=True,
             ),
         )
 

examples/storytelling-chatbot/src/bot_runner.py

@@ -217,7 +217,6 @@ if __name__ == "__main__":
     required_env_vars = [
         "GOOGLE_API_KEY",
         "DAILY_API_KEY",
-        "FAL_KEY",
         "ELEVENLABS_VOICE_ID",
         "ELEVENLABS_API_KEY",
     ]

examples/storytelling-chatbot/src/processors.py

@@ -1,16 +1,27 @@
+import os
 import re
 
+import google.ai.generativelanguage as glm
 from async_timeout import timeout
-from prompts import CUE_ASSISTANT_TURN, CUE_USER_TURN, IMAGE_GEN_PROMPT
+from loguru import logger
+from prompts import (
+    CUE_ASSISTANT_TURN,
+    CUE_USER_TURN,
+    FIRST_IMAGE_PROMPT,
+    IMAGE_GEN_PROMPT,
+    NEXT_IMAGE_PROMPT,
+)
 from utils.helpers import load_sounds
 
 from pipecat.frames.frames import (
     Frame,
     LLMFullResponseEndFrame,
+    LLMMessagesFrame,
     TextFrame,
     UserStoppedSpeakingFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.services.google import GoogleImageGenService, GoogleLLMContext, GoogleLLMService
 from pipecat.transports.services.daily import DailyTransportMessageFrame
 
 sounds = load_sounds(["talking.wav", "listening.wav", "ding.wav"])
@@ -44,24 +55,56 @@ class StoryImageProcessor(FrameProcessor):
     The processed frames are then yielded back.
 
     Attributes:
-        _fal_service (FALService): The FAL service, generates the images (fast fast!).
+        _image_gen_service: The FAL service, generates the images (fast fast!).
     """
 
-    def __init__(self, fal_service):
+    def __init__(self, image_gen_service):
         super().__init__()
-        self._fal_service = fal_service
+        self._image_gen_service = image_gen_service
+        # Create a new LLM service to use a different system prompt, etc
+        self._llm_service = GoogleLLMService(api_key=os.getenv("GOOGLE_API_KEY"))
+
+        self.pages = []
+        self.image_descriptions = []
+
+    def can_generate_metrics(self) -> bool:
+        return True
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         await super().process_frame(frame, direction)
 
-        if isinstance(frame, StoryImageFrame):
+        if isinstance(frame, StoryPageFrame):
+            # Special syntax for the first page
+            if self.pages == []:
+                prompt = FIRST_IMAGE_PROMPT % frame.text
+            else:
+                prompt = NEXT_IMAGE_PROMPT % (
+                    " ".join(self.pages),
+                    "; ".join(self.image_descriptions),
+                    frame.text,
+                )
+
+            await self.start_ttfb_metrics()
+            # TODO: This is coupled to google implementation now
+            txt = glm.Content(role="user", parts=[glm.Part(text=prompt)])
+            llm_response = await self._llm_service._client.generate_content_async(
+                contents=[txt], stream=False
+            )
+            image_description = llm_response.text
+            self.pages.append(frame.text)
+            self.image_descriptions.append(image_description)
             try:
-                async with timeout(7):
-                    async for i in self._fal_service.run_image_gen(IMAGE_GEN_PROMPT % frame.text):
+                async with timeout(15):
+                    async for i in self._image_gen_service.run_image_gen(
+                        IMAGE_GEN_PROMPT % image_description
+                    ):
                         await self.push_frame(i)
             except TimeoutError:
+                logger.debug("Image gen timeout")
                 pass
-            pass
+            await self.stop_ttfb_metrics()
+            # Push the StoryPageFrame so it gets TTS
+            await self.push_frame(frame)
         else:
             await self.push_frame(frame)
 
@@ -96,7 +139,8 @@ class StoryProcessor(FrameProcessor):
 
         elif isinstance(frame, TextFrame):
             # Add new text to the buffer
-            self._text += frame.text
+            # (character replace hack to fix TTS sequencing)
+            self._text += frame.text.replace(";", "—")
             # Process any complete patterns in the order they appear
             await self.process_text_content()
 

examples/storytelling-chatbot/src/prompts.py

@@ -4,35 +4,65 @@ LLM_BASE_PROMPT = {
         Your goal is to craft an engaging and fun story.
         Keep all responses short and no more than a few sentences.
         Start by asking the user what kind of story they'd like to hear. Don't provide any examples.
-        After they've answered the question, start telling the story. Include [break] after each sentence of the story.
+        After they've answered the question, start telling the story. Include three story sentences in your response. Add [break] after each sentence of the story.
         
-        Start each sentence with an image prompt, wrapped in triangle braces, that I can use to generate an illustration representing the upcoming scene.
+        EXAMPLE OUTPUT FORMAT:
+        story sentence 1 [break]
+        story sentence 2 [break]
+        story sentence 3 [break]
+        How would you like the story to continue?
+        END OF EXAMPLE OUTPUT
+
+        Generate three story sentences, then ask what should happen next and wait for my input. You can propose an idea for how the story should proceed, but make sure to tell me I can suggest whatever I want.
+        Please ensure your responses are less than 5 sentences long.
+        Please refrain from using any explicit language or content. Do not tell scary stories.
+        Once you've started telling the story, EVERY RESPONSE should follow the story sentence output format. It is VERY IMPORTANT that you continue to include [break] between story sentences. DO NOT RESPOND without story sentences and break tags.""",
+}
+
+
+IMAGE_GEN_PROMPT = "an illustration of %s. colorful, whimsical, painterly, concept art."
+
+CUE_USER_TURN = {"cue": "user_turn"}
+CUE_ASSISTANT_TURN = {"cue": "assistant_turn"}
+
+
+"""        Start each sentence with an image prompt, wrapped in triangle braces, that I can use to generate an illustration representing the upcoming scene.
         Image prompts should always be wrapped in triangle braces, like this: <image prompt goes here>.
         You should provide as much descriptive detail in your image prompt as you can to help recreate the current scene depicted by the sentence.
         For any recurring characters, you should provide a description of them in the image prompt each time, for example: <a brown fluffy dog ...>.
         Please do not include any character names in the image prompts, just their descriptions.
         Image prompts should focus on key visual attributes of all characters each time, for example <a brown fluffy dog and the tiny red cat ...>.
         Please use the following structure for your image prompts: characters, setting, action, and mood.
-        Image prompts should be less than 150-200 characters and start in lowercase.
+        Image prompts should be less than 150-200 characters and start in lowercase."""
 
-        STORY SENTENCE OUTPUT FORMAT:
-        <image description 1>
-        story sentence 1 [break]
-        <image description 2>
-        story sentence 2 [break]
-        <image description 3>
-        story sentence 3 [break]
-        How would you like the story to continue?
-        END OF EXAMPLE OUTPUT
+FIRST_IMAGE_PROMPT = """You are creating a prompt to generate an image for a child's story book.
 
-        Generate three story sentences, then ask what should happen next and wait for my input. You can propose an idea for how the story should proceed, but make sure to tell me I can suggest whatever I want. \
-        Please ensure your responses are less than 5 sentences long. \
-        Please refrain from using any explicit language or content. Do not tell scary stories.
-        Once you've started telling the story, EVERY RESPONSE should follow the story sentence output format. It is VERY IMPORTANT that you continue to include <image descriptions> and [break] between story sentences. DO NOT RESPOND without image descriptions and break tags.""",
-}
+You should provide as much descriptive detail in your image prompt as you can to help recreate the current scene depicted by the sentence.
+For any recurring characters, you should provide a description of them in the image prompt each time, for example: <a brown fluffy dog ...>.
+Please do not include any character names in the image prompts, just their descriptions.
+Image prompts should focus on key visual attributes of all characters each time, for example <a brown fluffy dog and the tiny red cat ...>.
+Please use the following structure for your image prompts: characters, setting, action, and mood.
+Image prompts should be less than 150-200 characters and start in lowercase.
 
 
-IMAGE_GEN_PROMPT = "illustrative art of %s. In the style of Studio Ghibli. colorful, whimsical, painterly, concept art."
+Here's the first page of the story:
+%s
+"""
 
-CUE_USER_TURN = {"cue": "user_turn"}
-CUE_ASSISTANT_TURN = {"cue": "assistant_turn"}
+NEXT_IMAGE_PROMPT = """You are creating a prompt to generate an image for a child's story book.
+
+Here is the text of the story so far:
+%s
+
+Here are the previous image prompts:
+%s
+
+You should provide as much descriptive detail in your image prompt as you can to help recreate the current scene depicted by the sentence.
+For any recurring characters, you should try to use the same description of them in the image prompt each time.
+Please do not include any character names in the image prompts, just their descriptions.
+Image prompts should focus on key visual attributes of all characters each time, for example <a brown fluffy dog and the tiny red cat ...>.
+Please use the following structure for your image prompts: characters, setting, action, and mood.
+Image prompts should be less than 150-200 characters and start in lowercase.
+Here's the next page of the story:
+%s
+"""