import torch
from diffusers import StableDiffusionImg2ImgPipeline
import cv2
import numpy as np
import gradio as gr
from PIL import Image

# Charger le modèle Stable Diffusion pour GPU ou CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe = pipe.to(device)

def generate_video_from_image(image, prompt, num_frames=10):
    """
    Génère une vidéo à partir d'une image de départ et d'un prompt.

    Args:
    - image (PIL.Image): Image de départ.
    - prompt (str): Texte descriptif pour générer des variations.
    - num_frames (int): Nombre de frames dans la vidéo.

    Returns:
    - str: Chemin du fichier vidéo généré.
    """
    frames = []
    image = image.resize((512, 512))  # Redimensionner l'image

    # Générer les frames
    for i in range(num_frames):
        generator = torch.Generator(device).manual_seed(42 + i)
        prompt_variation = f"{prompt}, variation {i}"
        generated_image = pipe(prompt=prompt_variation, image=image, generator=generator).images[0]
        frames.append(generated_image)

    # Dimensions de la vidéo
    width, height = frames[0].size
    video_path = "generated_video.mp4"

    # Créer un fichier vidéo
    video = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), 24, (width, height))
    for frame in frames:
        video.write(cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR))

    video.release()

    return video_path

# Interface utilisateur Gradio
css = """
h1#title {
    font-size: 3em;
    font-weight: bold;
    background: linear-gradient(45deg, #ff0000, #ff9900, #33cc33, #3399ff, #cc33ff);
    -webkit-background-clip: text;
    color: transparent;
    text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
    font-family: Arial, sans-serif;
    text-align: center;
}
"""

iface = gr.Interface(
    fn=generate_video_from_image,
    inputs=[
        gr.Image(type="pil", label="Image de départ"),
        gr.Textbox(label="Prompt", placeholder="Décris la scène ou l'effet désiré (ex: 'effet artistique, coucher de soleil')"),
        gr.Slider(1, 30, value=10, step=1, label="Nombre de Frames")
    ],
    outputs=gr.File(label="Vidéo générée"),
    title="Stable-FX05",
    css=css
)

iface.launch()