# -*- coding: utf-8 -*-
"""
VÍDEO — a família LTX-2 da Lightricks (LTX-2, LTX-2.3, LTX-2.5), via diffusers.

Pedido do dono: "quero um modelo de vídeo muito bom". Este é o caminho da
qualidade; o Wan 2.2 continua sendo o caminho do custo zero, e qual dos dois
responde é o `VIDEO_MODEL` — o `escolha.py` manda pela FAMÍLIA do id.

── O QUE O LEVANTAMENTO MOSTROU (Hub, 2026-09-15) ──────────────────────────
Perguntei ao Hub quem lidera image-to-video em diffusers, em vez de escolher
de memória. A ordem é: `Lightricks/LTX-2.5` (1,58 milhão de downloads, 3.932
curtidas), `Lightricks/LTX-2.3` (1,12 milhão), `Lightricks/LTX-Video` (762
mil). O Wan 2.2 TI2V-5B que está publicado aqui aparece com 172 mil. Não é
empate: a família LTX-2 é a escolha da comunidade por uma ordem de grandeza.

E ela NÃO é de graça, e este é o ponto que decide tudo:

    LTX-2.5        200,9 GB
    LTX-2.3        156,0 GB
    LTX-2.3-fp8     58,7 GB   <- o único que chega perto de caber
    Wan2.2-TI2V-5B  34,2 GB   <- o padrão de hoje

Num Space ZeroGPU gratuito o disco é efêmero: cada vez que ele dorme, os
pesos VOLTAM a ser baixados. 34 GB já são minutos; 200 GB não é lentidão, é
um recurso que nunca termina de subir. Por isso o padrão do registro CONTINUA
sendo o Wan — trocá-lo por decisão minha entregaria um vídeo que não gera.

── ENTÃO POR QUE ESTE ARQUIVO EXISTE AGORA ─────────────────────────────────
Porque antes ele era um ESQUELETO que levantava `NotImplementedError`, e o
`escolha.py` passou a rotear a família LTX para cá: apontar o `VIDEO_MODEL`
para um LTX daria um erro seco em vez de um vídeo. Um caminho que o roteamento
alcança tem de funcionar.

Com uma GPU dedicada (ver a tabela de custo em `docs/ZEROGPU.md`), o dono
troca UMA variável — `VIDEO_MODEL=Lightricks/LTX-2.3-fp8` — e o vídeo bom
entra sem deploy nenhum. É para isso que o registro existe.

── O QUE NÃO FOI VALIDADO, E É HONESTO DIZER ───────────────────────────────
A inferência NÃO rodou: não há GPU nem pesos neste ambiente, e baixar 58 GB
para provar uma chamada não caberia aqui de qualquer jeito. O que está
provado é o que dá para provar: que os ids existem, que o `model_index.json`
do LTX-2 declara `LTX2Pipeline`, e o roteamento por família
(`t_modelo_imagem.py`). O que conferir no primeiro deploy com GPU está
nomeado em `docs/MODELS.md`.
"""
from __future__ import annotations

import base64 as b64
import io
import logging
import os
import tempfile
from typing import Dict, Optional

from . import base

logger = logging.getLogger("worker.ltx")

VRAM_MINIMA_GB = 20

_pipe = None
_modelo_carregado = ""


def _modelo() -> str:
    return os.getenv("VIDEO_MODEL", "Lightricks/LTX-2.3-fp8").strip()


def baixar(progresso: base.Progresso) -> None:
    base.baixar_repo(_modelo(), progresso)


def carregar(progresso: base.Progresso) -> None:
    global _pipe, _modelo_carregado
    if base.modo_mock():
        _pipe = "mock"
        return
    if _pipe is not None and _modelo_carregado == _modelo():
        return
    progresso(None, "Carregando o modelo de vídeo…")
    import torch
    from diffusers import DiffusionPipeline
    # `DiffusionPipeline` resolve a classe pelo `model_index.json` do próprio
    # repositório (`LTX2Pipeline` no LTX-2). Importar a classe por nome
    # quebraria com ImportError num diffusers que a renomeasse, e prenderia
    # este arquivo a UMA versão da família — o oposto do que ele existe para
    # fazer.
    _pipe = DiffusionPipeline.from_pretrained(
        _modelo(), torch_dtype=torch.bfloat16,
        token=os.getenv("HF_TOKEN") or None)
    # Vídeo é o maior dos quatro: offload SEMPRE, mesmo com VRAM folgada — o
    # pico do VAE de vídeo no decode é traiçoeiro e é onde o OOM mora. É a
    # mesma decisão do wan22.py, e pelo mesmo motivo medido.
    _pipe.enable_model_cpu_offload()
    _modelo_carregado = _modelo()


def gerar(params: Dict, progresso: base.Progresso) -> base.Saida:
    if base.modo_mock():
        return b"", "video/mp4"

    prompt = (params.get("prompt") or "").strip()
    estilo = (params.get("style") or "").strip()
    if estilo:
        prompt = f"{prompt}, {estilo} style"
    largura = int(params.get("width") or 1280)
    altura = int(params.get("height") or 704)
    segundos = max(2, min(int(params.get("duration") or 5),
                          int(os.getenv("VIDEO_MAX_SEGUNDOS", "8"))))
    fps = int(params.get("fps") or 24)
    # O LTX quer contagem de quadros em 8n+1 (o VAE temporal comprime de 8 em
    # 8, e um resto sobrando sai como um quadro final corrompido). A conta
    # ARREDONDA PARA BAIXO: passar do que foi pedido custaria prazo de GPU num
    # modelo que já é o mais caro dos quatro.
    quadros = max(9, ((segundos * fps) // 8) * 8 + 1)
    passos = int(params.get("steps") or 30)

    imagem = None
    if params.get("image_url"):
        imagem = _abrir_imagem(params["image_url"], largura, altura)

    def a_cada_passo(pipe, passo, t, kw):
        progresso(int(95 * (passo + 1) / max(1, passos)), "Gerando o vídeo…")
        return kw

    extras = {"image": imagem} if imagem is not None else {}
    saida = _pipe(prompt=prompt,
                  negative_prompt=params.get("negative_prompt") or None,
                  width=largura, height=altura, num_frames=quadros,
                  frame_rate=fps,
                  num_inference_steps=passos,
                  generator=base.semente(params),
                  callback_on_step_end=a_cada_passo, **extras)
    return _para_mp4(saida.frames[0], fps), "video/mp4"


def descarregar() -> None:
    global _pipe
    _pipe = None
    base.limpar_vram()


def config_menor(params: Dict) -> Optional[Dict]:
    """Depois de um OOM: menos pixels, depois menos tempo, depois menos passos.

    A ORDEM é a do wan22.py e não é arbitrária — no vídeo o pico de memória é
    quadros × pixels, então cortar resolução rende mais que cortar passos, e
    cortar passos é o que mais estraga a imagem.
    """
    p = dict(params)
    w = int(p.get("width") or 1280)
    if w > 704:
        p["width"] = 704
        p["height"] = max(384, int((p.get("height") or 704) * 704 // max(1, w)))
        return p
    seg = int(p.get("duration") or 5)
    if seg > 3:
        p["duration"] = 3
        return p
    passos = int(p.get("steps") or 30)
    if passos > 20:
        p["steps"] = 20
        return p
    return None


def _abrir_imagem(dado: str, largura: int, altura: int):
    """A primeira imagem do image-to-video. Aceita data URL e endereço, que
    são as duas formas em que ela chega do app."""
    from PIL import Image
    if dado.startswith("data:"):
        bruto = b64.b64decode(dado.split(",", 1)[1])
    else:
        import requests
        bruto = requests.get(dado, timeout=30).content
    img = Image.open(io.BytesIO(bruto)).convert("RGB")
    return img.resize((largura, altura))


def _para_mp4(quadros, fps: int) -> bytes:
    """Os quadros viram um MP4 em arquivo temporário.

    `export_to_video` do diffusers escreve em disco e não devolve bytes, e o
    contrato desta casa é bytes — quem guarda é o `armazenamento.py`, que
    pode estar num bucket. O temporário morre aqui.
    """
    from diffusers.utils import export_to_video
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
        caminho = f.name
    try:
        export_to_video(quadros, caminho, fps=fps)
        with open(caminho, "rb") as f:
            return f.read()
    finally:
        try:
            os.unlink(caminho)
        except OSError:
            pass