from shiny import App, ui, reactive
from PIL import Image
import tempfile
import base64
import torch
from diffusers import StableDiffusionPipeline

# -----------------------------
# Charger le modèle vidéo
# -----------------------------
model_name = "Lightricks/ltx-video-distilled"
pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
pipe.to("cuda")  # ou "cpu" si tu n'as pas de GPU

# -----------------------------
# Interface utilisateur
# -----------------------------
app_ui = ui.page_fluid(
    ui.layout_sidebar(
        ui.panel_sidebar(
            ui.input_file("input_image", "Upload your image", accept=[".png", ".jpg", ".jpeg"]),
            ui.input_text("model_prompt", "Prompt", placeholder="Décris la vidéo que tu veux générer..."),
            ui.input_action_button("generate_video", "Générer la vidéo"),
        ),
        ui.panel_main(
            ui.output_ui("video_output")
        )
    )
)

# -----------------------------
# Backend pour générer la vidéo
# -----------------------------
def generate_video(img_file, prompt):
    if img_file is None or prompt.strip() == "":
        return ui.div("Merci de fournir une image et un prompt.")

    # Charger l'image
    img = Image.open(img_file[0]["datapath"])

    # Générer la vidéo
    video_result = pipe(prompt=prompt, init_image=img, num_inference_steps=50)
    video = video_result.videos[0]

    # Sauvegarder temporairement
    tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    video.save(tmp_file.name)

    # Encoder en base64 pour affichage HTML
    video_b64 = base64.b64encode(open(tmp_file.name, "rb").read()).decode()
    html_video = f'<video controls src="data:video/mp4;base64,{video_b64}" width="480"></video>'

    return ui.HTML(html_video)

# -----------------------------
# Reactive
# -----------------------------
def server(input, output, session):
    @reactive.event(input.generate_video)
    @output
    @ui.render_ui
    def video_output():
        return generate_video(input.input_image(), input.model_prompt())

# -----------------------------
# App Shiny
# -----------------------------
app = App(app_ui, server)