import gradio as gr
import torch
import cv2
import numpy as np
import sys

# Отладочная информация
print(f"Python: {sys.version}")
print(f"Torch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

from diffusers import DiffusionPipeline

# Загрузка модели БЕЗ .to("cuda") — accelerate сам управляет устройствами
pipe = DiffusionPipeline.from_pretrained(
    "damo-vilab/text-to-video-ms-1.7b",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    use_safetensors=True
)

# Оптимизации памяти (работают даже без явного .to("cuda"))
pipe.enable_vae_slicing()
if torch.cuda.is_available():
    pipe.enable_model_cpu_offload()

def generate_video(prompt):
    try:
        print(f"Generating: '{prompt}'")
        
        # Минимальные параметры для стабильности на T4
        video_frames = pipe(
            prompt,
            num_inference_steps=15,
            num_frames=16,
            guidance_scale=7.5
        ).frames[0]
        
        # Сохранение в /tmp
        output_path = "/tmp/output.mp4"
        
        frames_uint8 = [(frame * 255).astype(np.uint8) for frame in video_frames]
        height, width = frames_uint8[0].shape[:2]
        
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, 8, (width, height))
        
        for frame in frames_uint8:
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            out.write(frame_bgr)
        out.release()
        
        return output_path
    
    except Exception as e:
        print(f"ERROR: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# МИНИМАЛЬНЫЙ ИНТЕРФЕЙС БЕЗ УСТАРЕВШИХ ПАРАМЕТРОВ
demo = gr.Interface(
    fn=generate_video,
    inputs=gr.Textbox(
        label="Prompt (English)",
        value="a cat walking in a garden, cartoon style",
        lines=2
    ),
    outputs=gr.Video(label="Generated Video (16 frames)"),
    title="🎥 Text-to-Video Generator",
    description="Free via Hugging Face T4 GPU • Model: ModelScope",
    examples=[
        ["a robot dancing in cyberpunk city"],
        ["a panda eating bamboo in forest"],
        ["a bouncing ball on white background"]
    ],
    cache_examples=False  # КРИТИЧЕСКИ ВАЖНО для экономии памяти
)

if __name__ == "__main__":
    demo.launch()