from diffusers import StableDiffusionPipeline
from PIL import Image
import imageio
import gradio as gr
import torch

# Dynamic device and data type allocation
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# Load Stable Diffusion model
model_id = "runwayml/stable-diffusion-v1-5"  # Replace with your desired model ID
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, safety_checker=None)
pipe.to(device)

def generate_frames(prompt, num_frames=5, style="realistic"):
    """
    Generate coherent frames for animation.
    """
    frames = []
    for i in range(num_frames):
        frame_prompt = f"{prompt}, frame {i+1}, {style}"
        image = pipe(frame_prompt).images[0]
        frame_path = f"frame_{i+1}.png"
        image.save(frame_path)
        frames.append(frame_path)
    return frames

def create_animation(frames, output_path="output.gif", duration=500):
    """
    Combine frames into an animated GIF.
    """
    images = [Image.open(frame) for frame in frames]
    images[0].save(output_path, save_all=True, append_images=images[1:], duration=duration, loop=0)
    return output_path

def generate_gif(prompt, num_frames, style):
    """
    Generate an animation GIF from user input.
    """
    frames = generate_frames(prompt, num_frames, style)
    gif_path = create_animation(frames)
    return gif_path

# Gradio Interface
interface = gr.Interface(
    fn=generate_gif,
    inputs=[
        gr.Textbox(label="Story or Prompt", placeholder="Enter your short story or prompt here..."),
        gr.Slider(2, 10, value=5, label="Number of Frames"),
        gr.Textbox(label="Style", placeholder="e.g., realistic, cartoon, surreal")
    ],
    outputs=gr.File(label="Download Your Animation"),
    title="Text-to-Animation Generator",
    description="Generate short animations (GIFs) from a story or prompt using Stable Diffusion."
)

# Run the app
if __name__ == "__main__":
    interface.launch()