import gradio as gr
import torch
import numpy as np
import sys
from PIL import Image
from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
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,
    variant="fp16"
)

# Загрузка базовой модели (КАЧЕСТВЕННАЯ)
pipe = AnimateDiffPipeline.from_pretrained(
    "emilianJR/epiCRealism",
    motion_adapter=adapter,
    torch_dtype=torch.float16,
    variant="fp16"
)

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

# Перемещение на GPU
if torch.cuda.is_available():
    pipe = pipe.to("cuda")
    print("✓ Model loaded on GPU")

# Настройка планировщика
pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config,
    timestep_spacing="trailing"
)

def generate_video(image, prompt, negative_prompt="blurry, low quality, jittery"):
    try:
        # Конвертация изображения
        if image is not None and 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 (512×512 recommended)", type="numpy"),
        gr.Textbox(label="Prompt (describe motion)", value="gentle breeze blowing through hair, soft cinematic movement"),
        gr.Textbox(label="Negative Prompt", value="blurry, low quality, jittery motion, flickering")
    ],
    outputs=gr.Video(label="Generated Video (512×512, 16 frames, 2 seconds)"),
    title="🎥 High-Quality Image-to-Video Generator",
    description="✅ Apache 2.0 license — sell videos legally • Real img2vid support",
    examples=[
        [None, "slow cinematic camera pan, film grain", "blurry, jittery"],
        [None, "gentle breeze blowing through hair, soft movement", "low quality, flickering"],
        [None, "subtle floating particles, dreamy atmosphere", "jittery motion, artifacts"]
    ],
    cache_examples=False
)

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