from fastvideo import VideoGenerator
from fastvideo.v1.configs.sample.base import SamplingParam
from gtts import gTTS
import os
import torch

MODEL_PATH = "FastVideo/FastHunyuan-diffusers"
OUTPUT_DIR = "outputs"
generator = VideoGenerator.from_pretrained(model_path=MODEL_PATH, num_gpus=0)

default_params = SamplingParam.from_pretrained(MODEL_PATH)

def generate_video_with_audio(prompt, song_text=None, seed=None):
    params = default_params.copy()
    params.prompt = prompt
    params.num_frames = 30
    params.height = 512
    params.width = 512
    params.num_inference_steps = 25
    if seed is not None:
        params.seed = seed
    generator.generate_video(prompt=prompt, sampling_param=params)
    video_path = os.path.join(OUTPUT_DIR, f"{prompt[:30]}_{params.seed}.mp4")

    if song_text:
        tts = gTTS(song_text, lang="en")
        audio_path = os.path.join(OUTPUT_DIR, f"{prompt[:30]}_song.mp3")
        tts.save(audio_path)
        return video_path, audio_path
    return video_path, None