Update app.py

app.py

@@ -1,27 +1,62 @@
+import gradio as gr
 import requests
-import base64
 import io
-from PIL import Image
-import gradio as gr
 import os
+from PIL import Image
 
-API_KEY = os.environ["API_KEY"] # replace with your own API Token here
 API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
-HEADERS = {'Authorization': f'Bearer {API_KEY}'}
+api_key = os.environ.get('API_KEY')
+headers = {"Authorization": f"Bearer {api_key}"}
+
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+
+def generate_image_from_prompt(prompt_text):
+    image_bytes = query({"inputs": prompt_text})
+    generated_image = Image.open(io.BytesIO(image_bytes))
+    return generated_image
+
+title = "KVIImager 2.0 Demo 🎨"
+description = "This app uses Hugging Face AI model to generate an image based on the provided text prompt 🖼️."
+examples = [
+    ["A peaceful garden with a small cottage 🏡"],
+    ["A colorful abstract painting with geometric shapes 🎨"],
+    ["A serene beach at sunset 🌅"]
+]
+
+css_styles = {
+    'body': {
+        'background-color': '#f4f4f4',
+        'font-family': 'Arial, sans-serif'
+    },
+    'title': {
+        'color': 'navy',
+        'font-size': '36px',
+        'text-align': 'center',
+        'margin-bottom': '20px'
+    },
+    'textbox': {
+        'border': '2px solid #008CBA',
+        'border-radius': '5px',
+        'padding': '10px',
+        'margin-bottom': '20px',
+        'width': '300px',
+        'font-size': '16px'
+    },
+    'output_image': {
+        'box-shadow': '2px 2px 5px #888888'
+    }
+}
 
-def predictor(prompt):
-    payload = {'inputs': prompt}
-    
-    try:
-        response = requests.request("POST", url=API_URL, data=payload, headers=HEADERS)
-        
-        if not (response is None or response.status_code != 200):
-            result = response.text[7:-5]
-            
-            return Image.open(io.BytesIO(base64.b64decode(result)))
-                
-    except Exception as e:
-        print('Error processing request')
-        raise ValueError(e)
+input_prompt = gr.Textbox(label="Enter Prompt 📝", placeholder="E.g. 'A peaceful garden with a small cottage'")
+output_generated_image = gr.Image(label="Generated Image")
 
-iface = gr.Interface(fn=predictor, inputs="textbox", outputs='image').launch()
+gr.Interface(
+    generate_image_from_prompt, 
+    inputs=input_prompt, 
+    outputs=output_generated_image, 
+    title=title,
+    description=description,
+    examples=examples
+).launch(inline=True, css=css_styles)