Update app.py

app.py

@@ -1,13 +1,28 @@
+#------------------------------------------------------------------------
+# Import
+#------------------------------------------------------------------------
+
 import streamlit as st
 import requests
 from PIL import Image
 import io
 import os
 
-token = os.getenv("HF_TOKEN")
+#------------------------------------------------------------------------
+# HF API
+#------------------------------------------------------------------------
+
+# Retrieve the HF API key from environment variables 
+hf_api_key = os.getenv('HF_API_KEY')
+if not hf_api_key:
+    raise ValueError("HF_API_KEY not set in environment variables")
 
 API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-headers = {"Authorization": f"Bearer {token}"}
+headers = {"Authorization": f"Bearer {hf_api_key}"}
+
+#------------------------------------------------------------------------
+# Define functions
+#------------------------------------------------------------------------
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
@@ -35,5 +50,9 @@ def main():
         else:
             st.warning("Please enter a prompt.")
 
+#------------------------------------------------------------------------
+# Main Guard
+#------------------------------------------------------------------------
+
 if __name__ == "__main__":
     main()