import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import imageio
import numpy as np
from PIL import Image
import tempfile
import random
import os

# Lataa malli CPU:lle (muokkaa "cuda" jos GPU käytössä)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.to(device)
print("Model loaded!")

def generate_image(prompt, num_images=1):
    generator = torch.Generator(device=device).manual_seed(random.randint(0, 2**32-1))
    images = pipe(prompt=prompt, num_inference_steps=25, generator=generator).images
    return images[0]

def generate_video(prompt, duration=2, fps=5):
    frames = int(duration * fps)
    temp_dir = tempfile.mkdtemp()
    video_path = os.path.join(temp_dir, f"output_{random.randint(10000,99999)}.mp4")
    
    writer = imageio.get_writer(video_path, fps=fps)
    for i in range(frames):
        img = generate_image(prompt)
        img_array = np.array(img)
        writer.append_data(img_array)
    writer.close()
    return video_path

with gr.Blocks() as demo:
    gr.Markdown("# Simple SD Video Generator")
    prompt_input = gr.Textbox(label="Prompt", value="A majestic dragon flying over a castle")
    duration_input = gr.Slider(label="Duration (seconds)", minimum=1, maximum=10, value=2, step=1)
    output_video = gr.Video(label="Generated Video")
    
    gen_button = gr.Button("Generate Video")
    gen_button.click(fn=generate_video, inputs=[prompt_input, duration_input], outputs=[output_video])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)