import gradio as gr
import torch
import numpy as np
import sys
from PIL import Image
from diffusers import AnimateDiffPipeline, MotionAdapter
from diffusers.utils import export_to_video

# Отладка (выполняется ДО загрузки модели)
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)}")

# Загрузка адаптера движения
adapter = MotionAdapter.from_pretrained(
    "guoyww/animatediff-motion-adapter-v1-5-2",
    torch_dtype=torch.float16
)

# Загрузка базовой модели БЕЗ .to("cuda") и БЕЗ cpu_offload
pipe = AnimateDiffPipeline.from_pretrained(
    "emilianJR/epiCRealism",
    motion_adapter=adapter,
    torch_dtype=torch.float16
)

# ЕДИНСТВЕННАЯ оптимизация — безопасная для HF Spaces
pipe.enable_vae_slicing()

# Автоматическое определение устройства (работает в HF среде)
if torch.cuda.is_available():
    pipe = pipe.to("cuda")
    print("✓ Model moved to CUDA")
else:
    print("⚠️ Running on CPU (slow)")

def generate_video(image, prompt, negative_prompt="blurry, low quality"):
    try:
        # Конвертация изображения
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert("RGB")
        
        # Генерация видео
        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_frames=16,
            guidance_scale=7.5,
            num_inference_steps=25,
            generator=torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(42)
        )
        
        # Сохранение
        output_path = "/tmp/output.mp4"
        export_to_video(output.frames[0], output_path, fps=8)
        
        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.Image(label="Upload Image", type="numpy"),
        gr.Textbox(label="Prompt (describe motion)", value="gentle breeze blowing through hair"),
        gr.Textbox(label="Negative Prompt", value="blurry, low quality")
    ],
    outputs=gr.Video(label="Generated Video"),
    title="🎥 Commercial-Safe Image-to-Video",
    description="✅ Apache 2.0 license — sell videos legally",
    cache_examples=False
)

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