import gradio as gr
import numpy as np
from PIL import Image
import torch
import safetensors.torch as st
from diffusers import DiffusionPipeline

model_id = "./ckpts/snckrsgen.safetensors"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safe_tensors=True)
pipe.to("cuda")

# Function to generate an image from text using diffusion
def generate_image(prompt):
    
    images = pipe(prompt=prompt)
    return images[0]

# Gradio interface
inputs = gr.inputs.Textbox(lines=5, label="Enter text to generate image:")
outputs = gr.outputs.Image(label="Generated Image")

title = "ShoeGen: Generate an Image of a Shoe"
description = "Enter a text description of a shoe to generate an image of the shoe."
examples = [["A red shoe with white laces and black sole."], ["A blue sneaker with a white stripe."], ["A brown boot with a buckle."]]

gr.Interface(
    fn=generate_image,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    examples=examples
).launch()