from diffusers import StableDiffusionXLPipeline
import torch
from transformers.pipelines.image_to_text import Image
import gradio as gr

pipe = StableDiffusionXLPipeline.from_pretrained(
    "segmind/SSD-1B",
    torch_dtype = torch.float32,
    use_safetensors=True,
    # for cuda -- variant = "fp16",
)
#pipe.to("cuda")

prompt = "astronaut riding a green horse"
neg_prompt = "ugly, blurry, poor quality"

def generate_image(prompt, neg_prompt):
    image = pipe(
        prompt=prompt,
        negative_prompt=neg_prompt
    ).images[0]
    return image

prompt = gr.Text(
    label="Prompt",
    show_label=False,
    max_lines=1,
    placeholder="Enter your prompt :",
    container=False
)

neg_prompt = gr.Text(
    label="Negative Prompt",
    show_label=False,
    max_lines=1,
    placeholder="Enter your negative prompt",
    container=False
)

iface = gr.Interface(
     fn=generate_image,
     inputs=[prompt, neg_prompt],
     outputs="image",
     title="Text to Image Generation",
     examples=[
         ["a painting of a cute cat sitting on a chair","ugly, blurry"],
         ["an assortment riding a horse on mars","poorly down"]
     ],
     allow_flagging=False
)

iface.launch(share=True)
#image = pipe(prompt=prompt, negative_prompt = neg_prompt).images[0]
# print(image)