import gradio as gr
import torch
import spaces  # Импортируем модуль для работы с ZeroGPU
from diffusers import DiffusionPipeline

# 1. Функция генерации обязательно размечается декоратором @spaces.GPU
# По умолчанию резервируется 60 секунд. Если модель "тяжелая", можно указать duration (например, @spaces.GPU(duration=120))
@spaces.GPU
def generate_video(prompt, negative_prompt, progress=gr.Progress(track_tqdm=True)):
    
    # 2. Модель должна загружаться и переноситься на "cuda" строго ВНУТРИ этой функции.
    # За пределами этой функции вызывать .to("cuda") нельзя — ZeroGPU выдаст ошибку.
    pipe = DiffusionPipeline.from_pretrained(
        "adamdad/videocrafterv2_diffusers", 
        torch_dtype=torch.bfloat16
    ).to("cuda")
    
    # Рекомендуется для экономии памяти и ускорения
    pipe.enable_model_cpu_offload() 
    
    # 3. Запуск генерации
    # (Настройки кадров и шагов лучше держать минимальными, чтобы уложиться в лимит времени ZeroGPU)
    video_frames = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=25, 
        num_frames=16,
        height=320,
        width=512
    ).frames[0]
    
    # Здесь должен быть ваш код сохранения массива кадров в .mp4 файл
    video_path = "output_video.mp4"
    # (код сохранения, например через opencv или imageio)...
    
    return video_path

# 4. Простой интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🎬 VideoCrafter2 на ZeroGPU")
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Ваш запрос (Prompt)", placeholder="An astronaut riding a horse on mars, ultra realistic...")
            neg_prompt = gr.Textbox(label="Негативный запрос", value="low quality, blurry, static")
            btn = gr.Button("Сгенерировать видео")
        with gr.Column():
            video_output = gr.Video(label="Результат")
            
    btn.click(fn=generate_video, inputs=[prompt_input, neg_prompt], outputs=video_output)

demo.queue().launch()