Update app.py

now create images too

app.py

@@ -1,17 +1,34 @@
 from openai import OpenAI
 import gradio as gr
-import os
+import requests
+import matplotlib.pyplot as plt
+import io
+from PIL import Image, UnidentifiedImageError
+import random
+import time
 
 
-# create model for story creation
+#------------------------------------------------------------------------------------------------
+
+
+api_key = os.environ["genai_stories"]
+# define LLM model for story creation, with OpenAI's format
 client = OpenAI(
             base_url="https://api-inference.huggingface.co/v1/",
-	        api_key=os.environ["genai_stories"]
+	        api_key=api_key
             )
-
 # the model to utilize
 model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
+# define model for image creation
+# API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+# API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+# API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo"
+
+bearer_key = "Bearer " + api_key
+headers = {"Authorization": bearer_key}
+
 
 #------------------------------------------------------------------------------------------------
 
@@ -22,19 +39,76 @@ def chat_with_llm(model, messages):
         model=model,
         messages=messages,
         max_tokens=2048,
-        temperature=0)
+        temperature=0.7,
+        top_p=0.98,
+        seed=random.randint(0, 100))
     return completion.choices[0].message.content
 
 
 #------------------------------------------------------------------------------------------------
 
 
-# Initialize the story
+# connect with image model and create images
+def request_to_image_model(story):
+	response = requests.post(API_URL, headers=headers, json=story)
+	return response.content
+
+# create the final image, adjusting the model's hyperparameters and converting to RGB 
+def create_image(prompt):
+    try:
+        image_bytes = request_to_image_model({
+            "inputs": prompt, 
+            "parameters": {
+                # "num_inference_steps": 50,  # Optional: controls generation quality
+                # "guidance_scale": 7.5,      # Optional: controls adherence to the prompt
+                "seed": random.randint(0, 100)                 # Set the seed to a random number
+                }
+            })
+    except (UnidentifiedImageError) as e:
+        time.sleep(60)
+        image_bytes = request_to_image_model({
+            "inputs": prompt, 
+            "parameters": {
+                # "num_inference_steps": 50,  # Optional: controls generation quality
+                # "guidance_scale": 7.5,      # Optional: controls adherence to the prompt
+                "seed": random.randint(0, 100)                 # Set the seed to a random number
+                }
+            })
+
+    image = Image.open(io.BytesIO(image_bytes))
+    # Convert the image to RGB if it's not
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+
+    return image
+
+
+#------------------------------------------------------------------------------------------------
+
+
+# create a prompt for the image model based on the character
+def create_character_prompt(story_type):
+    main_character = story_type.split(' ')[2]
+    if main_character == 'king':
+        character_prompt = " The " + main_character + " has a strong jawline and piercing eyes, dark hair, shining silver armor with simple decorations, red cape falling behind him, large sword."
+    elif main_character == 'alien':
+        character_prompt = " The " + main_character + " has a round, shiny head, big eyes glow bright blue, small and green with tiny arms, silver spacesuit covers its body."
+    elif main_character == 'rabbit':
+        character_prompt = " The " + main_character + " has soft white fur, big ears stand up tall, wears a tiny red bow tie, fluffy tail bounces as it hops."
+    elif main_character == 'princess':
+        character_prompt = " The " + main_character + " has long brown hair, sparkling blue eyes shine brightly, wears a flowing pink gown, small jeweled crown rests on her head."
+    return character_prompt
+
+
+#------------------------------------------------------------------------------------------------
+
+
+# Initialize the story based on the story_type selected by the player
 def generate_story(story_type):
     global model, messages
     messages = []
-    messages.append({"role": "system", "content": f"You are a story creator for kids. The player is a kid, so you must use simple words. The player selects this type of story: {story_type}, according to which you will be asked to create a story. Keep the plot interesting and concise, and write the story as if you talk to children, with simple words. If you are asked to create a story, a question and the answers, then you must create these and only these. Do not create anything else. If you are asked to create the ending of the story, then you must create this and only this. Do not create anything else. "})
-    messages.append({"role": "user", "content": f"You must create and return ONLY a story part, a question and 4 answers. To do that, you must explicitly follow these steps: 1) Create the first part of the story, according to the story_type from your system content, within 50 and 60 words, without switching to new line. Then you must switch line by adding this: '\n\n'. 2) Then create a question, within 10 and 20 words, on how to proceed the story from step 1. Then you must switch line by adding this: '\n\n'. 3) Then create the 4 potential answers for the question in step 2. The answers of every question must be given in the format: '1:... | 2:... | 3:... | 4:...'. Do not change this format and do not add any new lines between the answers. Every answer must be maximum 20 words. All answers must be separated from each other with '|'. Don't explicitly specify 'Story', 'Question', or 'Answer'. You must reply in this format: '[story from step 1]\n\n[question from step 2]\n\n[answers from step 3]'. Do not create any other stuff."})
+    messages.append({"role": "system", "content": f"You are a structured storytelling assistant. Your purpose is to generate stories, questions, and answers in a specific format. Always adhere strictly to the rules provided by the user. This is the topic that you must create a story about: story_type = {story_type}"})
+    messages.append({"role": "user", "content": f"You must create and return one story part, one question and four answers. To do that, you must explicitly follow these steps: Step 1) Create the initial part of the story, according to the story_type from your system content, within 50 and 60 words, without switching to new line. Then you must switch line by adding this: '\n\n'. Step 2) Create a question, within 10 and 20 words, on how to proceed the story from step 1. Then you must switch line by adding this: '\n\n'. Step 3) Create the 4 potential answers for the question in step 2. The answers of the question must be given in the format: '1:... | 2:... | 3:... | 4:...'. Do not change this format and do not add any new lines between the answers. Every answer must be maximum 20 words. All answers must be separated from each other with '|'. Now some general guidelines for your response: 1) Don't explicitly specify 'Story', 'Question', or 'Answer'. 2) You must ALWAYS reply in this format: '[story from step 1]\n\n[question from step 2]\n\n[answers from step 3]'. 3) Do not return any other stuff in your response. 4) Always change lines with '\n\n' between ther story part and question and between the question and the answers."})
     messages = [messages[0]] + [messages[-1]]
     response = chat_with_llm(model, messages)
     lines    = response.split("\n\n")
@@ -42,13 +116,17 @@ def generate_story(story_type):
     question = lines[1]  # The second to last line is the question
     answers  = [i.strip() for i in lines[2].split('|')]  # The last line contains answers separated by '|'
     messages.append({"role": "assistant", "content": "I am waiting for next command."})
-    return story, question, gr.Radio(choices=answers, interactive=True), gr.Radio(choices=[story_type], interactive=False), gr.Button(interactive=False)
+    # main_character = story_type.split(' ')[2]
+    character_prompt = create_character_prompt(story_type)
+    image_prompt = "Cartoon image, with bright colors and simple shapes, that describes this story: " + story + character_prompt
+    image = create_image(image_prompt)
+    return story, question, gr.Radio(choices=answers, interactive=True), gr.Radio(choices=[story_type], interactive=False), gr.Button(interactive=False), image
 
 
 # Continue the story based on what happened so far and the player's latest answer
-def continue_story(previous_story, selected_option):
+def continue_story(previous_story, selected_option, story_type):
     global model, messages
-    messages.append({"role": "user", "content": f"You must create and return a story part, a question and 4 answers. To do that, you must explicitly follow these steps: 1) Based on this story so far: '{previous_story} {selected_option}', continue the story and create the next part of the story within 50 and 60 words, without changing lines. You must provide ONLY the new part that you created. Then you must switch line by adding this: '\n\n'. 2) Then create a question, within 10 and 20 words, on how to proceed the story from step 1. Then you must switch line by adding this: '\n\n'. 3) Then create the 4 potential answers for the question in step 2. The answers of every question must be given in the format: '1:... | 2:... | 3:... | 4:...'. Do not change this format and do not add any new lines between the answers. Every answer must be maximum 20 words. All answers must be separated from each other with '|' and no other separator. Don't explicitly specify 'Story', 'Question', or 'Answer'. You must reply in this format: '[story from step 1]\n\n[question from step 2]\n\n[answers from step 3]'. Do not create any other stuff."})
+    messages.append({"role": "user", "content": f"You must create and return one story part, one question and four answers. To do that, you must explicitly follow these steps: Step 1) Based on this story so far: '{previous_story} {selected_option}', continue the story and create the next part of the story within 50 and 60 words, without changing lines. You must provide ONLY the new part that you created. Then you must switch line by adding this: '\n\n'. Step 2) Create a question, within 10 and 20 words, on how to proceed the story from step 1. Then you must switch line by adding this: '\n\n'. Step 3) Create the 4 potential answers for the question in step 2. The answers of the question must be given in the format: '1:... | 2:... | 3:... | 4:...'. Do not change this format and do not add any new lines between the answers. Every answer must be maximum 20 words. All answers must be separated from each other with '|'. Now some general guidelines for your response: 1) Don't explicitly specify 'Story', 'Question', or 'Answer'. 2) You must ALWAYS reply in this format: '[story from step 1]\n\n[question from step 2]\n\n[answers from step 3]'. 3) Do not return any other stuff in your response. 4) Always change lines with '\n\n' between story part and question that you generate. 5) Always change lines with '\n\n' between question and the answers that you generate"})
     messages = [messages[0]] + [messages[-1]]
     response = chat_with_llm(model, messages)
     lines    = response.split("\n\n")
@@ -57,19 +135,46 @@ def continue_story(previous_story, selected_option):
     answers    = [i.strip() for i in lines[2].split('|')]  # The last line contains answers separated by '|'
     messages.append({"role": "assistant", "content": "I am waiting for next command."})
     story = previous_story + '\n\n' + next_story
-    return story, question, gr.Radio(choices=answers, interactive=True)
-
-
-# End the story
-def end_story(previous_story, selected_option):
+    # image_prompt = "Cartoon image of: " + next_story + ". The image has bright colors and simple shapes, intended for children story."
+    main_character = story_type.split(' ')[2].lower()
+    character_prompt = create_character_prompt(story_type)
+    if main_character in next_story.lower():
+        image_prompt = "Cartoon image, with bright colors and simple shapes, that describes this story: " + story + character_prompt
+    else:
+        image_prompt = "Cartoon image, with bright colors and simple shapes, that describes this story: " + story
+    image = create_image(image_prompt)
+    return story, question, gr.Radio(choices=answers, interactive=True), image
+
+
+# End the story based on what happened so far and the player's latest answer
+def end_story(previous_story, selected_option, story_type):
     global model, messages
-    messages.append({"role": "user", "content": f"Please create an ending for this story: '{previous_story}' and considering the latest answer: '{selected_option}'. You must provide only the ending of the story in an exciting way!."})
+    messages.append({"role": "user", "content": f"You must create an ending for this story: '{previous_story}' and also considering the latest answer: '{selected_option}'. You must provide only the ending of the story in an exciting way. Do not return any other stuff in your response."})
     end_story = chat_with_llm(model, messages)
     # lines    = response.split("\n\n")
     # end_story = lines[0]  # Everything before the last two lines is the story
     messages.append({"role": "assistant", "content": "I ended the story successfully. Now I am not waiting for any more responses from the player."})
     story = previous_story + '\n\n' + end_story
-    return story
+    # image_prompt = "Cartoon image of " + end_story + ". The image has bright colors and simple shapes, intended for children story."
+    main_character = story_type.split(' ')[2].lower()
+    character_prompt = create_character_prompt(story_type)
+    if main_character in end_story.lower():
+        image_prompt = "Cartoon image, with bright colors and simple shapes, that describes this story: " + story + character_prompt
+    else:
+        image_prompt = "Cartoon image, with bright colors and simple shapes, that describes this story: " + story
+    image = create_image(image_prompt)
+    return story, image
+
+
+# Function to handle resetting the app
+# def reset_app():
+#     # Return empty values for all outputs
+#     # global messages
+#     # messages = []
+#     return "", "", "", gr.Radio(choices=["A knight who leads an army attacking a castle",
+#                                                                     "A kid alien who lands on earth and explores around", 
+#                                                                     "Animals who participate in song contest", 
+#                                                                     "A little princess who is saved by a prince"], interactive=True), gr.Button(interactive=True) 
 
 
 # initialize messages 
@@ -82,14 +187,16 @@ with gr.Blocks() as game_ui:
         # Left column for the story
         with gr.Column(scale=1):
             story = gr.Textbox(label="Story", interactive=False, lines=12)
+            story_image = gr.Image()  # height=256, width=256
 
         # Right column for the question, answers, and buttons
         with gr.Column(scale=1):
-            story_type = gr.Radio(label="What story to create?", choices=["A knight who leads an army attacking a castle",
-                                                                    "A kid alien who lands on earth and explores around", 
-                                                                    "Animals who participate in song contest", 
-                                                                    "A little princess who is saved by a prince"])
-            
+            story_type = gr.Radio(label="What story to create?", choices=["A fearless king who leads an army attacking a castle.",
+                                                                          "A kid alien who lands on earth and explores around.", 
+                                                                          "A joyful rabbit who participates in song contest.", 
+                                                                          "A beautiful princess who tries to find her way home."
+                                                                          ])
+                
             start_button = gr.Button("Start Game!")
 
             question = gr.Textbox(label="Question", interactive=False)
@@ -98,11 +205,15 @@ with gr.Blocks() as game_ui:
             submit_button = gr.Button("Submit your answer")
             end_button = gr.Button("End the story")
 
-    # what the button do
-    start_button.click(fn=generate_story,  inputs=[story_type],     outputs=[story, question, answers, story_type, start_button])
-    submit_button.click(fn=continue_story, inputs=[story, answers], outputs=[story, question, answers])
-    end_button.click(fn=end_story,         inputs=[story, answers], outputs=[story])
+            # reset_button = gr.Button("Reset and play again")
+
+    # what the buttons do
+    start_button.click(fn=generate_story,  inputs=[story_type],                 outputs=[story, question, answers, story_type, start_button, story_image])
+    submit_button.click(fn=continue_story, inputs=[story, answers, story_type], outputs=[story, question, answers, story_image])
+    end_button.click(fn=end_story,         inputs=[story, answers, story_type], outputs=[story, story_image])
+    # reset_button.click(fn=reset_app,       inputs=[], outputs=[story, question, answers, story_type, start_button])
 
 
 # Launch the Gradio interface
-game_ui.launch(share=True)
+game_ui.launch()
+# game_ui.launch(share=True)