flux client to replicate + words count

app.py

@@ -6,7 +6,7 @@ import json
 import re
 from typing import List, Literal, TypedDict, Tuple
 from transformers import AutoTokenizer
-from gradio_client import Client
+import replicate
 from openai import OpenAI
 import anthropic
 from groq import Groq
@@ -41,7 +41,7 @@ MODEL_CONFIG: dict[ModelType, ModelConfig] = {
     },
     "CLAUDE": {
         "client": anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY")),
-        "model": "claude-3-5-sonnet-20240620",
+        "model": "claude-3-5-sonnet-20241022",
         "max_context": 128000,
         "tokenizer": AutoTokenizer.from_pretrained("Xenova/claude-tokenizer")
     },
@@ -275,7 +275,8 @@ def __predict():
         yield f"{C.EXCEPTION_KEYWORD} | {e}"
 
 
-def __generateImage(prompt: str):
+def __generateImageHF(prompt: str):
+    from gradio_client import Client
     fluxClient = Client(
         "black-forest-labs/FLUX.1-schnell",
         os.environ.get("HF_FLUX_CLIENT_TOKEN")
@@ -293,6 +294,28 @@ def __generateImage(prompt: str):
     return result
 
 
+def __generateImage(prompt: str):
+    fluxClient = replicate.Client(api_token=os.environ.get("REPLICATE_API_KEY"))
+    result = fluxClient.run(
+        "black-forest-labs/flux-schnell",
+        input={
+            "prompt": prompt,
+            "seed": 0,
+            "go_fast": False,
+            "megapixels": "1",
+            "num_outputs": 1,
+            "aspect_ratio": "4:3",
+            "output_format": "webp",
+            "output_quality": 80,
+            "num_inference_steps": 4
+        }
+    )
+    if result:
+        result = result[0]
+        U.pprint(f"imageResult={result}")
+        return result
+
+
 def __paintImageIfApplicable(
     imageContainer: DeltaGenerator,
     prompt: str,
@@ -312,7 +335,7 @@ def __paintImageIfApplicable(
                 unsafe_allow_html=True
             )
             imgContainer.image(C.IMAGE_LOADER)
-            (imagePath, seed) = __generateImage(imagePrompt)
+            imagePath = __generateImage(imagePrompt)
             imageContainer.image(imagePath)
     except Exception as e:
         U.pprint(e)
@@ -335,6 +358,19 @@ def __showButtons(options: list):
         )
 
 
+def __showWordsCount(response: str):
+    wordsCount = len(response.split())
+    countClass = "crossed-limit" if wordsCount > C.WORDS_LIMIT else ""
+    st.markdown(
+        f"""
+            <div class="words-count {countClass}">
+            {wordsCount} words
+            </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+
 def __resetButtonState():
     st.session_state.buttonValue = ""
 
@@ -494,12 +530,14 @@ def mainApp():
                     elif action:
                         U.pprint(f"{action=}")
                         if action == "SHOW_STORY_DATABASE":
-                            time.sleep(0.5)
+                            time.sleep(1)
                             st.switch_page("pages/popular-stories.py")
                         # st.code(jsonStr, language="json")
                 except Exception as e:
                     U.pprint(e)
 
+            __showWordsCount(response)
+
         saveLatestActivity()
 
 

constants.py

@@ -1,6 +1,7 @@
 JSON_SEPARATOR = ">>>>"
 EXCEPTION_KEYWORD = "<<EXCEPTION>>"
 BOOKING_LINK = "https://calendly.com"
+WORDS_LIMIT = 300
 
 SYSTEM_MSG = f"""
 => Context:

requirements.txt

@@ -7,4 +7,4 @@ anthropic
 supabase
 descope
 cloudinary
-
+replicate

utils.py

@@ -125,6 +125,18 @@ def applyCommonStyles():
             margin-bottom: 1rem;
         }
 
+        div.words-count {
+            font-size: 0.7rem !important;
+            color: green;
+            margin-top: -1rem;
+            text-align: right;
+        }
+
+        div.words-count.crossed-limit {
+            color: red;
+            font-size: 1rem !important;
+        }
+
         </style>
         """,
         unsafe_allow_html=True