import os
import io
from PIL import Image
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
hf_api_key = os.environ['HF_API_KEY']

# Helper function
import requests, json

# API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
API_URL = "https://api-inference.huggingface.co/models/cloudqi/cqi_text_to_image_pt_v0"

#Text-to-image endpoint
def get_completion(inputs, parameters=None, ENDPOINT_URL=API_URL):
    headers = {
      "Authorization": f"Bearer {hf_api_key}",
      "Content-Type": "application/json"
    }   
    data = { "inputs": inputs }
    if parameters is not None:
        data.update({"parameters": parameters})
    response = requests.request("POST",ENDPOINT_URL,headers=headers,data=json.dumps(data))
    return response.content

import gradio as gr 

def generate(prompt):
    output = get_completion(prompt)
    result_image = Image.open(io.BytesIO(output))
    return result_image


def loadGUI():
    gr.close_all()
    demo = gr.Interface(fn=generate,
                        inputs=[gr.Textbox(label="Your prompt")],
                        outputs=[gr.Image(label="Result")],
                        title="Image Generation with Stable Diffusion",
                        description="Generate any image with Stable Diffusion",
                        allow_flagging="never",
                        examples=["the spirit of a tamagotchi wandering in the city of Vienna","a mecha robot in a favela"])

    demo.launch(share=True)

def main():
     loadGUI()
     
     
if __name__ == "__main__":
     main()