import gradio as gr
from diffusers import StableVideoDiffusionPipeline, StableDiffusionPipeline
from PIL import Image
from gtts import gTTS
from moviepy.editor import ImageSequenceClip, AudioFileClip
from moviepy.audio.AudioClip import CompositeAudioClip
import torch, os

pipe_img = StableDiffusionPipeline.from_pretrained(
    "dreamlike-art/dreamlike-photoreal-2.0", torch_dtype=torch.float16
).to("cuda")

pipe_vid = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
).to("cuda")

os.makedirs("bgm_samples", exist_ok=True)
os.system("wget -q -nc -O bgm_samples/calm.mp3 https://github.com/mekarpeles/sample-audio/raw/main/calming.mp3")

def generate(prompt, lang):
    image = pipe_img(prompt, height=512, width=896).images[0]
    image.save("input.png")
    frames = pipe_vid(Image.open("input.png")).frames
    clip = ImageSequenceClip(frames, fps=6)

    tts = gTTS(text=prompt, lang=lang)
    tts.save("voice.mp3")

    music = AudioFileClip("bgm_samples/calm.mp3").subclip(0, clip.duration)
    voice = AudioFileClip("voice.mp3").set_duration(clip.duration)
    audio = CompositeAudioClip([music.volumex(0.3), voice])
    clip = clip.set_audio(audio)

    clip.write_videofile("output.mp4", codec="libx264", audio_codec="aac")
    return "output.mp4"

demo = gr.Interface(
    fn=generate,
    inputs=[gr.Textbox(label="Prompt"), gr.Dropdown(["vi", "en", "ja", "ko"], label="Language", value="vi")],
    outputs=gr.Video(label="🎬 Video kết quả"),
    title="Prompt2Video",
    description="Tạo video AI từ prompt với hình ảnh, chuyển động, giọng đọc và nhạc nền."
)

demo.launch()