added loaders + image generation

app.py

@@ -5,7 +5,9 @@ import pytz
 import time
 import json
 import re
+from typing import List
 from transformers import AutoTokenizer
+from gradio_client import Client
 
 from dotenv import load_dotenv
 load_dotenv()
@@ -46,7 +48,7 @@ Keep options to less than 9
 # Tier 1: Story Creation
 You initiate the storytelling process through a series of engaging prompts:
 Story Origin:
-Asks users to choose between personal anecdotes or adapting a well-known tale (creating a story database here of well-known stories to choose from).
+Asks users to choose between personal anecdotes or adapting a well-known story (creating a story database here of well-known finctional stories to choose from).
 
 Story Use Case:
 Asks users to define the purpose of building a story (e.g., profile story, for social media content).
@@ -120,6 +122,8 @@ Note that the final story should include twist, turns and events that make it re
 
 USER_ICON = "man.png"
 AI_ICON = "Kommuneity.png"
+IMAGE_LOADER = "ripple.svg"
+TEXT_LOADER = "balls.svg"
 START_MSG = "I want to create a story 😊"
 
 st.set_page_config(
@@ -144,13 +148,53 @@ pprint("\n")
 
 
 def __isInvalidResponse(response: str):
+    # new line followed by small case char
     if len(re.findall(r'\n[a-z]', response)) > 3:
         return True
 
+    # lot of repeating words
+    if re.findall(r'\b(\w+)(\s+\1){2,}\b', response) == 0:
+        return True
+
+    # json response without json separator
     if ('\n{\n    "options"' in response) and (JSON_SEPARATOR not in response):
         return True
 
 
+def __matchingKeywordsCount(keywords: List[str], text: str):
+    return sum([
+        1 if keyword in text else 0
+        for keyword in keywords
+    ])
+
+
+def __isStringNumber(s: str) -> bool:
+    try:
+        float(s)
+        return True
+    except ValueError:
+        return False
+
+
+def __getImageGenerationPrompt(prompt: str, response: str):
+    responseLower = response.lower()
+    if (
+        __matchingKeywordsCount(
+            ["adapt", "profile", "social media", "purpose", "use case"],
+            responseLower
+        ) > 2
+        and not __isStringNumber(prompt)
+        and prompt.lower() in responseLower
+    ):
+        return f'a scene from (({prompt})). Include main character'
+
+    if __matchingKeywordsCount(
+        ["Tier 2", "Tier-2"],
+        response
+    ) > 0:
+        return f"photo of a scene from this text: {response}"
+
+
 def __resetButtonState():
     st.session_state["buttonValue"] = ""
 
@@ -171,7 +215,10 @@ if "startMsg" not in st.session_state:
 
 def predict(prompt):
     historyFormatted = [{"role": "system", "content": SYSTEM_MSG}]
-    historyFormatted.extend(st.session_state.messages)
+    historyFormatted.extend([
+        {"role": message["role"], "content": message["content"]}
+        for message in st.session_state.messages
+    ])
     historyFormatted.append({"role": "user", "content": prompt })
     contextSize = countTokens(str(historyFormatted))
     pprint(f"{contextSize=}")
@@ -192,18 +239,37 @@ def predict(prompt):
             yield chunkContent
 
 
+def generateImage(prompt: str):
+    pprint(f"imagePrompt={prompt}")
+    client = Client("black-forest-labs/FLUX.1-schnell")
+    result = client.predict(
+            prompt=prompt,
+            seed=0,
+            randomize_seed=True,
+            width=1152,
+            height=896,
+            num_inference_steps=4,
+            api_name="/infer"
+    )
+    pprint(f"imageResult={result}")
+    return result
+
+
 st.title("Kommuneity Story Creator 📖")
-if not st.session_state.startMsg:
+if not (st.session_state["buttonValue"] or st.session_state["startMsg"]):
     st.button(START_MSG, on_click=lambda: __setStartMsg(START_MSG))
 
 for message in st.session_state.messages:
     role = message["role"]
     content = message["content"]
+    imagePath = message.get("image")
     avatar = AI_ICON if role == "assistant" else USER_ICON
     with st.chat_message(role, avatar=avatar):
         st.markdown(content)
+        if imagePath:
+            st.image(imagePath)
 
-if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_state.startMsg):
+if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_state["startMsg"]):
     __resetButtonState()
     __setStartMsg("")
 
@@ -213,27 +279,30 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
     st.session_state.messages.append({"role": "user", "content": prompt })
 
     with st.chat_message("assistant", avatar=AI_ICON):
-        placeholder = st.empty()
+        responseContainer = st.empty()
 
-        def getResponse():
+        def __printAndGetResponse():
             response = ""
+            # responseContainer.markdown(".....")
+            responseContainer.image(TEXT_LOADER)
             responseGenerator = predict(prompt)
 
             for chunk in responseGenerator:
                 response += chunk
                 if __isInvalidResponse(response):
+                    pprint(f"{response=}")
                     return
 
                 if JSON_SEPARATOR not in response:
-                    placeholder.markdown(response)
+                    responseContainer.markdown(response)
 
             return response
 
-        response = getResponse()
+        response = __printAndGetResponse()
         while not response:
             pprint("Empty response. Retrying..")
             time.sleep(0.5)
-            response = getResponse()
+            response = __printAndGetResponse()
 
         pprint(f"{response=}")
 
@@ -242,9 +311,22 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
             pprint(f"Selected: {optionLabel}")
 
         responseParts = response.split(JSON_SEPARATOR)
+
+        jsonStr = None
         if len(responseParts) > 1:
             [response, jsonStr] = responseParts
 
+        imagePath = None
+        try:
+            imagePrompt = __getImageGenerationPrompt(prompt, response)
+            if imagePrompt:
+                imageContainer = st.empty().image(IMAGE_LOADER)
+                (imagePath, seed) = generateImage(imagePrompt)
+                imageContainer.image(imagePath)
+        except Exception as e:
+            pprint(e)
+
+        if jsonStr:
             try:
                 json.loads(jsonStr)
                 jsonObj = json.loads(jsonStr)
@@ -260,4 +342,8 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
             except Exception as e:
                 pprint(e)
 
-    st.session_state.messages.append({"role": "assistant", "content": response})
+        st.session_state.messages.append({
+            "role": "assistant",
+            "content": response,
+            "image": imagePath,
+        })

requirements.txt

@@ -1,3 +1,4 @@
 python-dotenv
 groq
-transformers
+transformers
+gradio_client