Update app.py

app.py

@@ -1,26 +1,26 @@
 import os
 import gradio as gr
-import requests
+from diffusers import DiffusionPipeline
 from PIL import Image
 import io
 
-API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
 api_key = os.getenv("API_KEY")
 
-def query(payload):
-    headers = {"Authorization": f"Bearer {api_key}"}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.content
+def generate_image_with_stylization(prompt, style):
+    input_prompt = prompt
+    generated_image = pipeline.diffuse_text(input_prompt, stylize=True, style_prompt=style)
 
-def image_generation(prompt):
-    input_prompt = {
-        "inputs": prompt,
-    }
-    image_bytes = query(input_prompt)
-    image = Image.open(io.BytesIO(image_bytes))
-    return image
+    img = Image.fromarray(generated_image)
+    return img
 
 title = "Stable Fiddusion XL"
-description = "This app generates an image based on the provided prompt using the Stable Diffusion XL model."
+description = "This app generates an image based on the provided prompt using the Stable Diffusion XL model, applying stylization."
 
-gr.Interface(fn=image_generation, inputs="text", outputs="image", title=title, description=description).launch()
+gr.Interface(
+    fn=generate_image_with_stylization, 
+    inputs=["text", gr.inputs.Textbox(label="Style")], 
+    outputs="image", 
+    title=title, 
+    description=description
+).launch()