Update app.py

app.py

@@ -9,18 +9,19 @@ API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-
 headers = {"Authorization": f"Bearer {os.getenv('HF')}"}
 
 # Function to query the Hugging Face API
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
+@st.cache_data(ttl=300)  # Cache the function's output for 300 seconds (5 minutes)
+def query_image(prompt):
+    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
     return response.content
 
 # Check for 'text' parameter in URL
-query_params = st.query_params
-prompt_from_url = query_params.get('text')
+query_params = st.experimental_get_query_params()
+prompt_from_url = query_params.get('text', None)
 
 if prompt_from_url:
     # If 'text' parameter is present, generate and display the image only
-    prompt = prompt_from_url
-    image_bytes = query({"inputs": prompt})
+    prompt = prompt_from_url[0]
+    image_bytes = query_image(prompt)
     image = Image.open(io.BytesIO(image_bytes))
     st.image(image, caption="Generated Image", use_container_width=True)
 
@@ -45,7 +46,7 @@ else:
     if st.button("Generate Image"):
         if prompt:
             st.write("Generating image... Please wait ⏳")
-            image_bytes = query({"inputs": prompt})
+            image_bytes = query_image(prompt)
 
             # Display Image
             image = Image.open(io.BytesIO(image_bytes))