import torch
import gradio as gr
from diffusers import DiffusionPipeline

# Load the pre-trained model from Hugging Face
pipe = DiffusionPipeline.from_pretrained("alexx-ai/GBL2", torch_dtype=torch.float16)

# Move the model to GPU if available
pipe.to("cuda")

# Function to generate an image based on the input prompt
def generate_image(prompt):
    # Generate the image using the model's pipeline
    image = pipe(prompt).images[0]
    return image

# Create the Gradio interface for text-to-image generation
with gr.Blocks() as demo:
    gr.Markdown("### Text to Image Generator using GBL2")
    
    # Input text box for the prompt
    prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Type something like 'A beautiful sunset in Ghibli style'")
    
    # Output image display
    image_output = gr.Image(label="Generated Image")
    
    # Set up the button to trigger the image generation
    prompt_input.submit(generate_image, inputs=prompt_input, outputs=image_output)

# Launch the Gradio app
demo.launch()