import os 
import gradio as gr
from huggingface_hub import InferenceClient


class Text_to_image:
  def __init__(self):
    HUGGINGFACE_API_TOKEN=os.getenv("HUGGINGFACE_API_TOKEN")
    self.client = InferenceClient(token=HUGGINGFACE_API_TOKEN,model="stabilityai/stable-diffusion-xl-base-1.0")

  def text_to_image(self,prompt):
    image = self.client.text_to_image(prompt)
    return image
  
  def gradio_interface(self):
    
    with gr.Blocks(theme='JohnSmith9982/small_and_pretty',css="style.css") as demo:
      gr.HTML("""
                <center><h1 style="color:#02C160">Text To Image</h1></center>""")
      gr.HTML("""
                <center><h6 style="color:#02C160">Wait few seconds after Regenerated</h6></center>""")
      with gr.Column(elem_id = "col-container"):
        output=gr.Image( )
      with gr.Row(elem_id = "col-container"):
          with gr.Column():
            prompt=gr.Textbox(show_label=False,placeholder="Enter your prompt")
          with gr.Column():
            button=gr.Button("Generate Image",elem_classes="button")
      button.click(self.text_to_image,prompt,output)
      prompt.submit(self.text_to_image,prompt,output)
    demo.launch(debug=True)
  

if __name__=="__main__":
  text_to_image=Text_to_image()
  text_to_image.gradio_interface()