import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image

# Load the model
pipe = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-1")
pipe.enable_model_cpu_offload()

# Define the input and output functions
def text_to_image(prompt):
    generator = torch.Generator(device="cpu").manual_seed(0)
    image = pipe(prompt, num_inference_steps=25, generator=generator).images[0]
    return image

# Create a placeholder
placeholder = "A photograph of the inside of a subway train. There are raccoons sitting on the seats. One of them is reading a newspaper. The window shows the city in the background."

# Create the Gradio interface
title = "Kandinsky 3.0"
description = "This model generates an image based on a given text prompt."
how_to_use = "Input a description of the image you want to generate, for example: 'A forest with a river and a bridge under the moonlight.'"
examples = [["A dark alley with flickering streetlights and a mysterious figure lurking in the shadows"], 
            ["A futuristic cityscape with neon lights and flying cars"]]

gr.Interface(fn=text_to_image, inputs=gr.Textbox(placeholder=placeholder), outputs=gr.Image()).launch()