import torch
import gradio as gr
from diffusers import DiffusionPipeline

# Загрузка модели
model_id = "ali-vilab/text-to-video-ms-1.7b"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe = pipe.to("cpu")  # Установка CPU как устройства для модели

def generate_video(prompt: str, num_inference_steps: int = 50, guidance_scale: float = 7.5):
  """
  Генерирует видео по текстовому описанию.

  Args:
      prompt (str): Текстовое описание видео.
      num_inference_steps (int, optional): Количество шагов вывода. По умолчанию 50.
      guidance_scale (float, optional): Масштаб направляющей. По умолчанию 7.5.

  Returns:
      str: Путь к сгенерированному видео.
  """

  video_frames = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).frames
  video_path = f"generated_video_{prompt[:10]}.mp4"
  video_frames[0].save(video_path, save_all=True, append_images=video_frames[1:])
  return video_path

with gr.Blocks() as demo:
  with gr.Row():
    with gr.Column():
      text_input = gr.Textbox(label="Введите текстовое описание видео:")
      num_inference_steps_slider = gr.Slider(minimum=10, maximum=100, value=50, step=10, label="Количество шагов вывода:")
      guidance_scale_slider = gr.Slider(minimum=1, maximum=15, value=7.5, step=0.5, label="Масштаб направляющей:")
      generate_button = gr.Button("Сгенерировать видео")
    with gr.Column():
      video_output = gr.Video(label="Сгенерированное видео:")

  generate_button.click(fn=generate_video, inputs=[text_input, num_inference_steps_slider, guidance_scale_slider], outputs=video_output)

demo.launch()