app.py

import gradio as gr
import torch
import soundfile as sf
import numpy as np
from diffusers import StableVideoDiffusionPipeline
from transformers import MusicgenForConditionalGeneration, MusicgenProcessor

# Определяем устройство (GPU, если доступен)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Загружаем модель для генерации видео (Stable Video Diffusion)
video_pipe = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

# Загружаем модель для генерации музыки (MusicGen)
music_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
music_processor = MusicgenProcessor.from_pretrained("facebook/musicgen-small")

# Функция генерации видео
def generate_video(prompt):
    # ❗ Stable Video Diffusion требует изображение, а не текст ❗
    init_image = video_pipe.text_to_image(prompt, num_inference_steps=25).images[0]  # Генерация стартового кадра
    video_frames = video_pipe(init_image, num_inference_steps=25).frames  # Генерация видео

    # Сохранение видео
    video_path = "generated_video.mp4"
    video_pipe.save_video(video_frames, video_path, fps=15)

    return video_path

# Функция генерации музыки
def generate_music(prompt):
    inputs = music_processor(text=[prompt], return_tensors="pt")  # Правильная передача текста
    audio_output = music_model.generate(**inputs)  # Генерация аудио

    # Декодирование аудиовыхода
    audio_array = audio_output[0].cpu().detach().numpy().squeeze()

    # Сохранение аудио в файл
    audio_path = "generated_audio.wav"
    sf.write(audio_path, audio_array, samplerate=16000)  # Правильная запись WAV

    return audio_path

# Интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("# OnlyReels.AI 🎥🎵 - Генерация видео с музыкой")

    with gr.Row():
        prompt_input = gr.Textbox(label="Введите описание видео и музыки", placeholder="Пример: Киберпанк-город ночью, синтезаторная музыка")

    with gr.Row():
        btn_video = gr.Button("🎬 Сгенерировать видео")
        btn_music = gr.Button("🎵 Сгенерировать музыку")

    video_output = gr.Video(label="Сгенерированное видео")
    music_output = gr.Audio(label="Сгенерированная музыка")

    btn_video.click(fn=generate_video, inputs=prompt_input, outputs=video_output)
    btn_music.click(fn=generate_music, inputs=prompt_input, outputs=music_output)

# Запуск интерфейса
demo.launch()