Update app.py

app.py

@@ -1,33 +1,29 @@
+import requests
+import base64
+import io
+from PIL import Image
 import gradio as gr
-from diffusers import DiffusionPipeline
 
-# Load the DiffusionPipeline and LORA weights
-pipeline = DiffusionPipeline.from_pretrained("stablediffusionapi/juggernaut-xl-v5")
-pipeline.load_lora_weights("Kvikontent/kviimager2.0")
+API_TOKEN = os.environ["API_TOKEN"] # replace with your own API Token here
+API_URL = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
+HEADERS = {
+    'Content-Type': 'application/json',
+    'Authorization': f'{API_TOKEN}'
+}
 
-# Define the function to handle user input and generate the image
-def generate_image(prompt):
-    image = pipeline(prompt)
-    return image
+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)
 
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs="text",
-    outputs="image",
-    title="KVIImager 2.0 Demo",
-    description="Enter a prompt to generate an image"
-)
-
-# Add examples for the user input
-examples = [
-    "Sunset over the ocean",
-    "Starry night sky"
-]
-
-iface.set_config(
-    examples=examples
-)
-
-# Launch the Gradio interface
-iface.launch()
+iface = gr.Interface(fn=predictor, inputs="textbox", outputs='image').launch()