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)

import gradio as gr 

def generate(prompt, negative_prompt, steps, guidance, width, height):
    params = {
        "negative_prompt": negative_prompt,
        "num_inference_steps": steps,
        "guidance_scale": guidance,
        "width": width,
        "height": height
    }
    
    output = get_completion(prompt, params)
    pil_image = Image.open(io.BytesIO(output))
    return pil_image

def loadGUI():
    gr.close_all()
    demo = gr.Interface(fn=generate,
                    inputs=[
                        gr.Textbox(label="Your prompt"),
                        gr.Textbox(label="Negative prompt"),
                        gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=25,
                                 info="In how many steps will the denoiser denoise the image?"),
                        gr.Slider(label="Guidance Scale", minimum=1, maximum=20, value=7, 
                                  info="Controls how much the text prompt influences the result"),
                        gr.Slider(label="Width", minimum=64, maximum=512, step=64, value=512),
                        gr.Slider(label="Height", minimum=64, maximum=512, step=64, value=512),
                    ],
                    outputs=[gr.Image(label="Result")],
                    title="Image Generation with Stable Diffusion",
                    description="Generate any image with Stable Diffusion",
                    allow_flagging="never"
                    )

    demo.launch(share=True)

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