import os
import time
import gradio as gr
import torch
from PIL import Image

from huggingface_hub import snapshot_download
from diffusers import StableVideoDiffusionPipeline
import imageio.v2 as imageio

REPO_ID = "stabilityai/stable-video-diffusion-img2vid-xt"

# Где хранить файлы модели (пытаемся /data, если нет — /home/user)
ROOT = "/data" if os.path.isdir("/data") else "/home/user"
MODEL_DIR = os.path.join(ROOT, "models", "svd-xt")
OUT_DIR = "/tmp/outputs"
os.makedirs(OUT_DIR, exist_ok=True)

pipe = None


def ensure_model():
    os.makedirs(MODEL_DIR, exist_ok=True)

    # Если уже скачано — не качаем заново
    if any(os.path.exists(os.path.join(MODEL_DIR, f)) for f in ["model_index.json", "config.json"]):
        print("✅ Model already present in:", MODEL_DIR)
        return

    print("⬇️ Downloading model to:", MODEL_DIR)
    snapshot_download(
        repo_id=REPO_ID,
        local_dir=MODEL_DIR,
        local_dir_use_symlinks=False,
        resume_download=True,
        max_workers=4,
        token=os.environ.get("HF_TOKEN"),  # если модель gated
    )
    print("✅ Download finished. Top files:", os.listdir(MODEL_DIR)[:30])


def get_pipe():
    global pipe
    if pipe is not None:
        return pipe

    ensure_model()

    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    pipe = StableVideoDiffusionPipeline.from_pretrained(
        MODEL_DIR,
        torch_dtype=dtype,
        local_files_only=True,  # <-- важно: после скачивания не лезем в интернет
    )

    if torch.cuda.is_available():
        pipe = pipe.to("cuda")
        try:
            pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass

    return pipe


def run(image: Image.Image, motion: int, fps: int, frames: int, steps: int, seed: int):
    if image is None:
        raise gr.Error("Upload an image first.")

    if seed < 0:
        seed = int(time.time()) % 10_000_000

    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))

    pipe = get_pipe()
    img = image.convert("RGB")

    out = pipe(
        image=img,
        num_frames=int(frames),
        fps=int(fps),
        motion_bucket_id=int(motion),
        num_inference_steps=int(steps),
        generator=generator,
    )

    video_frames = out.frames[0]  # list[PIL]
    out_path = os.path.join(OUT_DIR, f"svd_{int(time.time())}.mp4")
    imageio.mimsave(out_path, video_frames, fps=int(fps))
    return out_path


with gr.Blocks(title="SVD img2vid XT") as demo:
    gr.Markdown("## Stable Video Diffusion (img2vid-xt)\nModel downloads once at startup (resume enabled).")

    with gr.Row():
        inp = gr.Image(type="pil", label="Input image")
        out = gr.Video(label="Output video")

    with gr.Accordion("Settings", open=False):
        motion = gr.Slider(1, 255, value=100, step=1, label="motion_bucket_id (lower = calmer)")
        fps = gr.Slider(6, 30, value=12, step=1, label="fps")
        frames = gr.Slider(8, 30, value=14, step=1, label="num_frames")
        steps = gr.Slider(10, 50, value=25, step=1, label="steps")
        seed = gr.Number(value=-1, precision=0, label="seed (-1 random)")

    btn = gr.Button("Generate", variant="primary")
    btn.click(run, [inp, motion, fps, frames, steps, seed], out)

demo.queue().launch()