import requests
import base64
import io
from PIL import Image
import gradio as gr
import os

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}'}

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)

iface = gr.Interface(fn=predictor, inputs="textbox", outputs='image').launch()