from transformers import pipeline
import gradio as gr
from PIL import Image

# Carrega o pipeline de geração de imagens
generator = pipeline("image-generation", model="stabilityai/stable-diffusion-2")

def generate_frame_with_model(prompt, width=256, height=256):
    result = generator(prompt, height=height, width=width, num_inference_steps=50)
    image = Image.open(result[0]["generated_image"]).convert("RGB")
    return image

def create_animation_with_model(prompt, frame_count, duration):
    frames = [generate_frame_with_model(f"{prompt}, frame {i}") for i in range(frame_count)]
    gif_path = "output.gif"
    frames[0].save(gif_path, save_all=True, append_images=frames[1:], duration=duration, loop=0)
    return gif_path

def animate_with_model(prompt, frame_count, duration):
    gif_path = create_animation_with_model(prompt, frame_count, duration)
    return gif_path

# Interface Gradio
with gr.Blocks() as app:
    gr.Markdown("# Gerador de Animações com Prompt de Texto")
    prompt = gr.Textbox(label="Descrição do Sprite", placeholder="Descreva o sprite ou a animação que deseja criar")
    frame_count = gr.Slider(label="Número de Frames", minimum=1, maximum=10, value=5, step=1)
    duration = gr.Slider(label="Duração do Frame (ms)", minimum=50, maximum=500, value=100, step=10)
    animation_output = gr.Image(label="Animação Gerada")
    generate_button = gr.Button("Gerar Animação")

    generate_button.click(animate_with_model, inputs=[prompt, frame_count, duration], outputs=animation_output)

app.launch()