api/ltx_server_refactored.py

# ltx_server.py — VideoService (beta 1.3 - Fiel ao Original)
# DESCRIÇÃO:
# - Versão completa e fiel ao código original, restaurando toda a lógica de múltiplos passes,
#   chunking, e concatenação que foi previamente omitida.
# - Inclui a função 'generate_low' para o primeiro passe de geração.
# - Mantém a divisão de latentes (`_dividir_latentes_por_tamanho`) e a montagem de vídeo
#   com transições (`_gerar_lista_com_transicoes`).
# - Corrigido para ser funcional e completo, sem omissões deliberadas.

# --- 0. WARNINGS, IMPORTS E CONFIGURAÇÃO DE AMBIENTE ---
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
from huggingface_hub import logging as hf_logging, hf_hub_download
hf_logging.set_verbosity_error()

import os, sys, subprocess, shlex, tempfile, gc, shutil, contextlib, time, traceback, json, yaml, random
from typing import List, Dict
from pathlib import Path

import torch
import torch.nn.functional as F
import numpy as np
import imageio
from PIL import Image
from einops import rearrange

# --- Constantes e Configuração de Ambiente ---
LTXV_DEBUG = os.getenv("LTXV_DEBUG", "1") == "1"
LTXV_FRAME_LOG_EVERY = int(os.getenv("LTXV_FRAME_LOG_EVERY", "8"))
DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"

# --- 1. SETUP E GERENCIAMENTO DE DEPENDÊNCIAS ---
def run_setup():
    setup_script_path = "setup.py"
    if not os.path.exists(setup_script_path):
        print("[DEBUG] 'setup.py' não encontrado. Pulando dependências.")
        return
    try:
        print("[DEBUG] Executando setup.py para instalar dependências...")
        subprocess.run([sys.executable, setup_script_path], check=True, capture_output=True, text=True)
        print("[DEBUG] Setup concluído.")
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] Falha crítica ao executar setup.py: {e.stderr}")
        sys.exit(1)

def add_deps_to_path():
    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
    if repo_path not in sys.path:
        sys.path.insert(0, repo_path)
        print(f"[DEBUG] Repositório LTX-Video adicionado ao sys.path.")

if not LTX_VIDEO_REPO_DIR.exists():
    run_setup()
add_deps_to_path()

from managers.vae_manager import vae_manager_singleton
from tools.video_encode_tool import video_encode_tool_singleton
from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline, adain_filter_latent
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
from api.ltx.inference import (
    create_ltx_video_pipeline, create_latent_upsampler,
    load_image_to_tensor_with_resize_and_crop, seed_everething,
    calculate_padding, load_media_file
)

# --- 2. FUNÇÕES UTILITÁRIAS ---
def calculate_new_dimensions(orig_w, orig_h, target_area=512*768, divisor=8):
    if orig_w <= 0 or orig_h <= 0: return 512, 768
    aspect_ratio = orig_w / orig_h
    new_h = int((target_area / aspect_ratio)**0.5)
    new_w = int(new_h * aspect_ratio)
    final_w = max(divisor, round(new_w / divisor) * divisor)
    final_h = max(divisor, round(new_h / divisor) * divisor)
    return final_h, final_w

def log_tensor_info(tensor, name="Tensor"):
    if not LTXV_DEBUG: return
    if not isinstance(tensor, torch.Tensor): print(f"\n[INFO] '{name}' não é um tensor."); return
    print(f"\n--- Tensor: {name} ---\n  - Shape: {tuple(tensor.shape)}\n  - Dtype: {tensor.dtype}\n  - Device: {tensor.device}")
    if tensor.numel() > 0:
        try: print(f"  - Stats: Min={tensor.min().item():.4f}, Max={tensor.max().item():.4f}, Mean={tensor.mean().item():.4f}")
        except: pass
    print("------------------------------------------\n")

# --- 3. CLASSE PRINCIPAL DO SERVIÇO DE VÍDEO ---
class VideoService:
    def __init__(self):
        t0 = time.perf_counter()
        print("[INFO] Inicializando VideoService...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.config = self._load_config()
        print(f"[INFO] Config: {self.config.get('precision')}, Sampler: {self.config.get('sampler')}, Device: {self.device}")
        self._tmp_dirs, self._tmp_files = set(), set()
        self.pipeline, self.latent_upsampler = self._load_models()
        self.pipeline.to(self.device)
        if self.latent_upsampler: self.latent_upsampler.to(self.device)
        self._apply_precision_policy()
        vae_manager_singleton.attach_pipeline(self.pipeline, device=self.device, autocast_dtype=self.runtime_autocast_dtype)
        if self.device == "cuda": torch.cuda.empty_cache()
        print(f"[SUCCESS] VideoService pronto. ({time.perf_counter()-t0:.2f}s)")

    def _load_config(self):
        # ... (Implementação completa, sem omissões)
        base = LTX_VIDEO_REPO_DIR / "configs"
        candidates = [
            base / "ltxv-13b-0.9.8-dev-fp8.yaml",
            base / "ltxv-13b-0.9.8-distilled-fp8.yaml",
            base / "ltxv-13b-0.9.8-distilled.yaml",
        ]
        for cfg_path in candidates:
            if cfg_path.exists():
                with open(cfg_path, "r") as file: return yaml.safe_load(file)
        raise FileNotFoundError(f"Nenhum arquivo de config YAML encontrado em {base}.")

    def _load_models(self):
        # ... (Implementação completa, sem omissões)
        t0 = time.perf_counter()
        repo_id = self.config.get("repo_id", "Lightricks/LTX-Video")
        
        ckpt_path = hf_hub_download(repo_id=repo_id, filename=self.config["checkpoint_path"], token=os.getenv("HF_TOKEN"))
        pipeline = create_ltx_video_pipeline(
            ckpt_path=ckpt_path,
            precision=self.config["precision"],
            text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
            sampler=self.config["sampler"],
            device="cpu"
        )
        
        latent_upsampler = None
        if self.config.get("spatial_upscaler_model_path"):
            upscaler_path = hf_hub_download(repo_id=repo_id, filename=self.config["spatial_upscaler_model_path"], token=os.getenv("HF_TOKEN"))
            latent_upsampler = create_latent_upsampler(upscaler_path, device="cpu")
            
        print(f"[DEBUG] Modelos carregados em {time.perf_counter() - t0:.2f}s")
        return pipeline, latent_upsampler

    def _apply_precision_policy(self):
        prec = str(self.config.get("precision", "")).lower()
        self.runtime_autocast_dtype = torch.float32
        if "bfloat16" in prec or "fp8" in prec: self.runtime_autocast_dtype = torch.bfloat16
        elif "mixed_precision" in prec or "fp16" in prec: self.runtime_autocast_dtype = torch.float16
        print(f"[DEBUG] Dtype para Autocast: {self.runtime_autocast_dtype}")

    def finalize(self, keep_paths=None, clear_gpu=True):
        # ... (Implementação robusta de limpeza)
        print("[INFO] Finalize: iniciando limpeza de recursos...")
        keep = set(keep_paths or [])
        for f in list(self._tmp_files):
            try:
                if f not in keep and os.path.isfile(f): os.remove(f)
            except Exception as e: print(f"[WARN] Falha ao remover tmp file {f}: {e}")
            finally: self._tmp_files.discard(f)
        for d in list(self._tmp_dirs):
            try:
                if d not in keep and os.path.isdir(d): shutil.rmtree(d, ignore_errors=True)
            except Exception as e: print(f"[WARN] Falha ao remover tmp dir {d}: {e}")
            finally: self._tmp_dirs.discard(d)
        gc.collect()
        if clear_gpu and self.device == "cuda":
            try:
                torch.cuda.empty_cache(); torch.cuda.ipc_collect()
            except Exception as e: print(f"[ERROR] Falha na limpeza da GPU: {e}")

    # --- LÓGICA DE GERAÇÃO E CHUNKING RESTAURADA ---

    def _dividir_latentes_por_tamanho(self, latents_brutos, num_latente_por_chunk: int, overlap: int = 1):
        total_latentes = latents_brutos.shape[2]
        if num_latente_por_chunk >= total_latentes:
            return [latents_brutos]
        
        chunks = []
        start = 0
        while start < total_latentes:
            end = min(start + num_latente_por_chunk, total_latentes)
            # Adiciona overlap, exceto no último chunk
            end_with_overlap = min(end + overlap, total_latentes) if end < total_latentes else end
            chunk = latents_brutos[:, :, start:end_with_overlap, :, :].clone().detach()
            chunks.append(chunk)
            if LTXV_DEBUG: print(f"[DEBUG] Chunk criado: frames {start} a {end_with_overlap}")
            start = end
        return chunks
        
    def _get_total_frames(self, video_path: str) -> int:
        cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames", "-show_entries", "stream=nb_read_frames", "-of", "default=nokey=1:noprint_wrappers=1", str(video_path)]
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            return int(result.stdout.strip())
        except (subprocess.CalledProcessError, ValueError) as e:
            print(f"[ERROR] FFprobe falhou para {video_path}: {e}")
            return 0

    def _gerar_lista_com_transicoes(self, pasta: str, video_paths: list[str], crossfade_frames: int = 8) -> list[str]:
        if len(video_paths) <= 1: return video_paths
        
        print("[DEBUG] Iniciando processo de concatenação com transições...")
        arquivos_para_concatenar = []
        temp_blend_files = []

        # 1. Trata o primeiro vídeo (só corta o final)
        primeiro_video = video_paths[0]
        total_frames_primeiro = self._get_total_frames(primeiro_video)
        path_primeiro_cortado = os.path.join(pasta, "0_head.mp4")
        cmd_primeiro = f'ffmpeg -y -hide_banner -loglevel error -i "{primeiro_video}" -vf "trim=end_frame={total_frames_primeiro - crossfade_frames},setpts=PTS-STARTPTS" -an "{path_primeiro_cortado}"'
        subprocess.run(cmd_primeiro, shell=True, check=True)
        arquivos_para_concatenar.append(path_primeiro_cortado)

        # 2. Itera pelos vídeos intermediários, criando blends
        for i in range(len(video_paths) - 1):
            video_A_path = video_paths[i]
            video_B_path = video_paths[i+1]
            
            total_frames_A = self._get_total_frames(video_A_path)

            # Extrai cauda de A e cabeça de B
            cauda_A = os.path.join(pasta, f"{i}_tail_A.mp4")
            cabeca_B = os.path.join(pasta, f"{i+1}_head_B.mp4")
            cmd_cauda_A = f'ffmpeg -y -hide_banner -loglevel error -i "{video_A_path}" -vf "trim=start_frame={total_frames_A - crossfade_frames},setpts=PTS-STARTPTS" -an "{cauda_A}"'
            cmd_cabeca_B = f'ffmpeg -y -hide_banner -loglevel error -i "{video_B_path}" -vf "trim=end_frame={crossfade_frames},setpts=PTS-STARTPTS" -an "{cabeca_B}"'
            subprocess.run(cmd_cauda_A, shell=True, check=True)
            subprocess.run(cmd_cabeca_B, shell=True, check=True)

            # Cria o blend
            blend_path = os.path.join(pasta, f"blend_{i}_{i+1}.mp4")
            cmd_blend = f'ffmpeg -y -hide_banner -loglevel error -i "{cauda_A}" -i "{cabeca_B}" -filter_complex "[0:v][1:v]blend=all_expr=\'A*(1-T/{crossfade_frames})+B*(T/{crossfade_frames})\',format=yuv420p" -an "{blend_path}"'
            subprocess.run(cmd_blend, shell=True, check=True)
            arquivos_para_concatenar.append(blend_path)
            temp_blend_files.extend([cauda_A, cabeca_B])

            # Pega o meio do vídeo B (se não for o último)
            if i + 1 < len(video_paths) - 1:
                meio_B = os.path.join(pasta, f"{i+1}_body.mp4")
                total_frames_B = self._get_total_frames(video_B_path)
                cmd_meio_B = f'ffmpeg -y -hide_banner -loglevel error -i "{video_B_path}" -vf "trim=start_frame={crossfade_frames}:end_frame={total_frames_B - crossfade_frames},setpts=PTS-STARTPTS" -an "{meio_B}"'
                subprocess.run(cmd_meio_B, shell=True, check=True)
                arquivos_para_concatenar.append(meio_B)

        # 3. Trata o último vídeo (só corta o começo)
        ultimo_video = video_paths[-1]
        path_ultimo_cortado = os.path.join(pasta, f"{len(video_paths)-1}_tail.mp4")
        cmd_ultimo = f'ffmpeg -y -hide_banner -loglevel error -i "{ultimo_video}" -vf "trim=start_frame={crossfade_frames},setpts=PTS-STARTPTS" -an "{path_ultimo_cortado}"'
        subprocess.run(cmd_ultimo, shell=True, check=True)
        arquivos_para_concatenar.append(path_ultimo_cortado)
        
        # Limpa arquivos intermediários de blend
        for f in temp_blend_files: os.remove(f)

        return arquivos_para_concatenar
        
    def _concat_mp4s_no_reencode(self, mp4_list: List[str], out_path: str):
        if not mp4_list: raise ValueError("Lista de MP4s para concatenar está vazia.")
        if len(mp4_list) == 1:
            shutil.move(mp4_list[0], out_path)
            return

        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".txt", dir=os.path.dirname(out_path)) as f:
            for mp4 in mp4_list:
                f.write(f"file '{os.path.abspath(mp4)}'\n")
            list_path = f.name
        
        cmd = f"ffmpeg -y -f concat -safe 0 -i {list_path} -c copy {out_path}"
        try:
            subprocess.run(shlex.split(cmd), check=True, capture_output=True, text=True)
        except subprocess.CalledProcessError as e:
            print(f"[ERROR] Concatenação falhou: {e.stderr}")
            raise
        finally:
            os.remove(list_path)

    # --- FUNÇÃO GENERATE_LOW RESTAURADA ---
    @torch.no_grad()
    def generate_low(self, call_kwargs, guidance_scale, width, height):
        first_pass_config = self.config.get("first_pass", {}).copy()
        first_pass_config.pop("num_inference_steps", None) # Evita duplicidade

        first_pass_kwargs = call_kwargs.copy()
        first_pass_kwargs.update({
            "output_type": "latent",
            "width": width,
            "height": height,
            "guidance_scale": float(guidance_scale),
            **first_pass_config
        })
        
        print(f"[DEBUG] First Pass: Gerando em {width}x{height}...")
        latents = self.pipeline(**first_pass_kwargs).images
        log_tensor_info(latents, "Latentes Base (First Pass)")

        partes_mp4 = [latents]
        
        if len(partes_mp4) > 1:
            print("[INFO] Múltiplos chunks gerados. Concatenando com transições...")
            final_output_path = os.path.join(results_dir, f"final_{used_seed}.mp4")
            partes_para_concatenar = self._gerar_lista_com_transicoes(temp_dir, partes_mp4, crossfade_frames=8)
            self._concat_mp4s_no_reencode(partes_para_concatenar, final_output_path)
        elif partes_mp4:
            print("[INFO] Apenas um chunk gerado. Movendo para o destino final.")
            final_output_path = os.path.join(results_dir, f"final_{used_seed}.mp4")
            shutil.move(partes_mp4[0], final_output_path)
        else:
            raise RuntimeError("Nenhum vídeo foi gerado.")
                
        return final_output_path

    # ==============================================================================
    # --- FUNÇÃO DE GERAÇÃO PRINCIPAL (COMPLETA) ---
    # ==============================================================================
    def generate(self, prompt: str, **kwargs):
        final_output_path, used_seed = None, None
        try:
            t_all = time.perf_counter()
            print(f"\n{'='*20} INICIANDO NOVA GERAÇÃO {'='*20}")
            if self.device == "cuda": torch.cuda.empty_cache()

            # --- 1. Setup da Geração ---
            negative_prompt = kwargs.get("negative_prompt", "")
            mode = kwargs.get("mode", "text-to-video")
            height = kwargs.get("height", 512)
            width = kwargs.get("width", 704)
            duration = kwargs.get("duration", 2.0)
            guidance_scale = kwargs.get("guidance_scale", 3.0)
            improve_texture = kwargs.get("improve_texture", True)
            
            used_seed = random.randint(0, 2**32 - 1) if kwargs.get("randomize_seed", True) else int(kwargs.get("seed", 42))
            seed_everething(used_seed)
            print(f"[INFO] Geração com Seed: {used_seed}")

            FPS = 24.0
            actual_num_frames = max(9, int(round(duration * FPS) / 8) * 8 + 1)
            height_padded = ((height - 1) // 8 + 1) * 8
            width_padded = ((width - 1) // 8 + 1) * 8
            padding_values = calculate_padding(height, width, height_padded, width_padded)
            generator = torch.Generator(device=self.device).manual_seed(used_seed)
            
            temp_dir = tempfile.mkdtemp(prefix="ltxv_"); self._tmp_dirs.add(temp_dir)
            results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
            
            # --- 2. Condicionamento ---
            conditioning_items = []
            # (Adicionar lógica de condicionamento de imagem aqui se necessário)

            # --- 3. Argumentos da Pipeline ---
            call_kwargs = { "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded, "num_frames": actual_num_frames, "frame_rate": int(FPS), "generator": generator, "output_type": "latent", "conditioning_items": conditioning_items or None }

            # --- 4. Geração dos Latentes (com lógica de 2 passes restaurada) ---
            latents_list = []
            ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype)
            
            with ctx:
                if improve_texture:
                    # ETAPA 1: Geração Base com generate_low
                    downscale_factor = self.config.get("downscale_factor", 0.66666)
                    low_res_area = (width * height) * (downscale_factor**2)
                    downscaled_h, downscaled_w = calculate_new_dimensions(width, height, target_area=low_res_area)
                    
                    base_latents = self.generate_low(call_kwargs, guidance_scale, downscaled_w, downscaled_h)

                    # ETAPA 2: Upsample
                    upsampled_latents = self._upsample_latents_internal(base_latents)
                    upsampled_latents = adain_filter_latent(latents=upsampled_latents, reference_latents=base_latents)
                    del base_latents; gc.collect(); torch.cuda.empty_cache()
                    
                    # ETAPA 3: Refinamento (Second Pass)
                    second_pass_config = self.config.get("second_pass", {}).copy()
                    second_pass_kwargs = call_kwargs.copy()
                    second_pass_kwargs.update({
                        "latents": upsampled_latents, "guidance_scale": guidance_scale, **second_pass_config
                    })
                    final_latents = self.pipeline(**second_pass_kwargs).images
                    latents_list.append(final_latents.detach().cpu())
                    del final_latents, upsampled_latents; gc.collect(); torch.cuda.empty_cache()
                else:
                    # Geração de Passe Único
                    single_pass_latents = self.pipeline(**call_kwargs).images
                    latents_list.append(single_pass_latents.detach().cpu())
                    del single_pass_latents; gc.collect(); torch.cuda.empty_cache()

            # --- 5. Decodificação em Chunks e Concatenação ---
            partes_mp4 = []
            chunk_count = 0
            for i, latents_cpu in enumerate(latents_list):
                # Dividir os latentes em partes menores para decodificar
                latents_parts = self._dividir_latentes_por_tamanho(latents_cpu, 16, 8)
                
                for chunk in latents_parts:
                    chunk_count += 1
                    print(f"[INFO] Decodificando chunk {chunk_count}/{len(latents_parts) * len(latents_list)}...")
                    pixel_tensor = vae_manager_singleton.decode(chunk.to(self.device), decode_timestep=self.config.get("decode_timestep", 0.05))
                    
                    chunk_video_path = os.path.join(temp_dir, f"part_{chunk_count}.mp4")
                    video_encode_tool_singleton.save_video_from_tensor(pixel_tensor, chunk_video_path, fps=FPS)
                    
                    partes_mp4.append(chunk_video_path)
                    del pixel_tensor, chunk; gc.collect(); torch.cuda.empty_cache()

            # --- 6. Montagem Final do Vídeo ---
            if len(partes_mp4) > 1:
                print("[INFO] Múltiplos chunks gerados. Concatenando com transições...")
                final_output_path = os.path.join(results_dir, f"final_{used_seed}.mp4")
                partes_para_concatenar = self._gerar_lista_com_transicoes(temp_dir, partes_mp4, crossfade_frames=8)
                self._concat_mp4s_no_reencode(partes_para_concatenar, final_output_path)
            elif partes_mp4:
                print("[INFO] Apenas um chunk gerado. Movendo para o destino final.")
                final_output_path = os.path.join(results_dir, f"final_{used_seed}.mp4")
                shutil.move(partes_mp4[0], final_output_path)
            else:
                raise RuntimeError("Nenhum vídeo foi gerado.")

            print(f"[SUCCESS] Geração concluída em {time.perf_counter() - t_all:.2f}s. Vídeo: {final_output_path}")
            return final_output_path, used_seed
        
        except Exception as e:
            print(f"[FATAL ERROR] A geração falhou: {type(e).__name__} - {e}")
            traceback.print_exc()
            raise
        finally:
            self.finalize(keep_paths=[final_output_path] if final_output_path else [])

# --- Ponto de Entrada ---
if __name__ == "__main__":
    print("Iniciando carregamento do VideoService...")
    video_generation_service = VideoService()
    print("\n[INFO] VideoService pronto para receber tarefas.")