api/ltx_server.py

# ltx_server.py — VideoService (beta 1.1)
# Sempre output_type="latent"; no final: VAE (bloco inteiro) → pixels → MP4.
# Ignora UserWarning/FutureWarning e injeta VAE no manager com dtype/device corretos.

# --- 0. WARNINGS E AMBIENTE ---
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*")

from huggingface_hub import logging, hf_hub_download

logging.set_verbosity_error()
logging.set_verbosity_warning()
logging.set_verbosity_info()
logging.set_verbosity_debug()
LTXV_DEBUG=1
LTXV_FRAME_LOG_EVERY=8

# --- 1. IMPORTAÇÕES ---
import os, subprocess, shlex, tempfile
import torch
import json
import numpy as np
import random
import os
import shlex
import yaml
from typing import List, Dict
from pathlib import Path
import imageio
import tempfile
from huggingface_hub import hf_hub_download
import sys
import subprocess
import gc
import shutil
import contextlib
import time
import traceback
from einops import rearrange
import torch.nn.functional as F

# Singletons (versões simples)
from managers.vae_manager import vae_manager_singleton
from tools.video_encode_tool import video_encode_tool_singleton

# --- 2. GERENCIAMENTO DE DEPENDÊNCIAS E SETUP ---
def _query_gpu_processes_via_nvml(device_index: int) -> List[Dict]:
    try:
        import psutil
        import pynvml as nvml
        nvml.nvmlInit()
        handle = nvml.nvmlDeviceGetHandleByIndex(device_index)
        try:
            procs = nvml.nvmlDeviceGetComputeRunningProcesses_v3(handle)
        except Exception:
            procs = nvml.nvmlDeviceGetComputeRunningProcesses(handle)
        results = []
        for p in procs:
            pid = int(p.pid)
            used_mb = None
            try:
                if getattr(p, "usedGpuMemory", None) is not None and p.usedGpuMemory not in (0,):
                    used_mb = max(0, int(p.usedGpuMemory) // (1024 * 1024))
            except Exception:
                used_mb = None
            name = "unknown"
            user = "unknown"
            try:
                import psutil
                pr = psutil.Process(pid)
                name = pr.name()
                user = pr.username()
            except Exception:
                pass
            results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
        nvml.nvmlShutdown()
        return results
    except Exception:
        return []

def _query_gpu_processes_via_nvidiasmi(device_index: int) -> List[Dict]:
    cmd = f"nvidia-smi -i {device_index} --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits"
    try:
        out = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT, text=True, timeout=2.0)
    except Exception:
        return []
    results = []
    for line in out.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 3:
            try:
                pid = int(parts[0]); name = parts[1]; used_mb = int(parts[2])
                user = "unknown"
                try:
                    import psutil
                    pr = psutil.Process(pid)
                    user = pr.username()
                except Exception:
                    pass
                results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
            except Exception:
                continue
    return results

def calculate_new_dimensions(orig_w, orig_h, divisor=8):
    """
    Calcula novas dimensões mantendo a proporção, garantindo que ambos os
    lados sejam divisíveis pelo divisor especificado (padrão 8).
    """
    if orig_w == 0 or orig_h == 0:
        # Retorna um valor padrão seguro
        return 512, 512

    # Preserva a orientação (paisagem vs. retrato)
    if orig_w >= orig_h:
        # Paisagem ou quadrado
        aspect_ratio = orig_w / orig_h
        # Começa com uma altura base e calcula a largura
        new_h = 512 # Altura base para paisagem
        new_w = new_h * aspect_ratio
    else:
        # Retrato
        aspect_ratio = orig_h / orig_w
        # Começa com uma largura base e calcula a altura
        new_w = 512 # Largura base para retrato
        new_h = new_w * aspect_ratio
        
    # Arredonda AMBOS os valores para o múltiplo mais próximo do divisor
    final_w = int(round(new_w / divisor)) * divisor
    final_h = int(round(new_h / divisor)) * divisor

    # Garante que as dimensões não sejam zero após o arredondamento
    final_w = max(divisor, final_w)
    final_h = max(divisor, final_h)

    print(f"[Dimension Calc] Original: {orig_w}x{orig_h} -> Calculado: {new_w:.0f}x{new_h:.0f} -> Final (divisível por {divisor}): {final_w}x{final_h}")
    return final_h, final_w # Retorna (altura, largura)

def handle_media_upload_for_dims(filepath, current_h, current_w):
    """
    Esta função agora usará o novo cálculo robusto.
    (O corpo desta função não precisa de alterações, pois ela já chama a função de cálculo)
    """
    if not filepath or not os.path.exists(str(filepath)):
        return gr.update(value=current_h), gr.update(value=current_w)
    try:
        if str(filepath).lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
            with Image.open(filepath) as img:
                orig_w, orig_h = img.size
        else: # Assumir que é um vídeo
            with imageio.get_reader(filepath) as reader:
                meta = reader.get_meta_data()
                orig_w, orig_h = meta.get('size', (current_w, current_h))
        
        # Chama a nova função corrigida
        new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
        
        return gr.update(value=new_h), gr.update(value=new_w)
    except Exception as e:
        print(f"Erro ao processar mídia para dimensões: {e}")
        return gr.update(value=current_h), gr.update(value=current_w)


def _gpu_process_table(processes: List[Dict], current_pid: int) -> str:
    if not processes:
        return "  - Processos ativos: (nenhum)\n"
    processes = sorted(processes, key=lambda x: (x.get("used_mb") or 0), reverse=True)
    lines = ["  - Processos ativos (PID | USER | NAME | VRAM MB):"]
    for p in processes:
        star = "*" if p["pid"] == current_pid else " "
        used_str = str(p["used_mb"]) if p.get("used_mb") is not None else "N/A"
        lines.append(f"    {star} {p['pid']} | {p['user']} | {p['name']} | {used_str}")
    return "\n".join(lines) + "\n"

def run_setup():
    setup_script_path = "setup.py"
    if not os.path.exists(setup_script_path):
        print("[DEBUG] 'setup.py' não encontrado. Pulando clonagem de dependências.")
        return
    try:
        print("[DEBUG] Executando setup.py para dependências...")
        subprocess.run([sys.executable, setup_script_path], check=True)
        print("[DEBUG] Setup concluído com sucesso.")
    except subprocess.CalledProcessError as e:
        print(f"[DEBUG] ERRO no setup.py (code {e.returncode}). Abortando.")
        sys.exit(1)

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,
)

DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
if not LTX_VIDEO_REPO_DIR.exists():
    print(f"[DEBUG] Repositório não encontrado em {LTX_VIDEO_REPO_DIR}. Rodando setup...")
    run_setup()

def add_deps_to_path():
    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
    if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
        sys.path.insert(0, repo_path)
        print(f"[DEBUG] Repo adicionado ao sys.path: {repo_path}")

add_deps_to_path()

# --- 3. IMPORTAÇÕES ESPECÍFICAS DO MODELO ---
from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
from ltx_video.pipelines.pipeline_ltx_video import adain_filter_latent


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

# --- 5. CLASSE PRINCIPAL DO SERVIÇO ---
class VideoService:
    def __init__(self):
        t0 = time.perf_counter()
        print("[DEBUG] Inicializando VideoService...")
        self.debug = os.getenv("LTXV_DEBUG", "1") == "1"
        self.frame_log_every = int(os.getenv("LTXV_FRAME_LOG_EVERY", "8"))
        self.config = self._load_config()
        print(f"[DEBUG] Config carregada (precision={self.config.get('precision')}, sampler={self.config.get('sampler')})")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[DEBUG] Device selecionado: {self.device}")
        self.last_memory_reserved_mb = 0.0
        self._tmp_dirs = set(); self._tmp_files = set(); self._last_outputs = []

        self.pipeline, self.latent_upsampler = self._load_models()
        print(f"[DEBUG] Pipeline e Upsampler carregados. Upsampler ativo? {bool(self.latent_upsampler)}")

        print(f"[DEBUG] Movendo modelos para {self.device}...")
        self.pipeline.to(self.device)
        if self.latent_upsampler:
            self.latent_upsampler.to(self.device)

        self._apply_precision_policy()
        print(f"[DEBUG] runtime_autocast_dtype = {getattr(self, 'runtime_autocast_dtype', None)}")

        # Injeta pipeline/vae no manager (impede vae=None)
        vae_manager_singleton.attach_pipeline(
            self.pipeline,
            device=self.device,
            autocast_dtype=self.runtime_autocast_dtype
        )
        print(f"[DEBUG] VAE manager conectado: has_vae={hasattr(self.pipeline, 'vae')} device={self.device}")

        if self.device == "cuda":
            torch.cuda.empty_cache()
            self._log_gpu_memory("Após carregar modelos")

        print(f"[DEBUG] VideoService pronto. boot_time={time.perf_counter()-t0:.3f}s")

    def _log_gpu_memory(self, stage_name: str):
        if self.device != "cuda":
            return
        device_index = torch.cuda.current_device() if torch.cuda.is_available() else 0
        current_reserved_b = torch.cuda.memory_reserved(device_index)
        current_reserved_mb = current_reserved_b / (1024 ** 2)
        total_memory_b = torch.cuda.get_device_properties(device_index).total_memory
        total_memory_mb = total_memory_b / (1024 ** 2)
        peak_reserved_mb = torch.cuda.max_memory_reserved(device_index) / (1024 ** 2)
        delta_mb = current_reserved_mb - getattr(self, "last_memory_reserved_mb", 0.0)
        processes = _query_gpu_processes_via_nvml(device_index) or _query_gpu_processes_via_nvidiasmi(device_index)
        print(f"\n--- [LOG GPU] {stage_name} (cuda:{device_index}) ---")
        print(f"  - Reservado: {current_reserved_mb:.2f} MB / {total_memory_mb:.2f} MB  (Δ={delta_mb:+.2f} MB)")
        if peak_reserved_mb > getattr(self, "last_memory_reserved_mb", 0.0):
            print(f"  - Pico reservado (nesta fase): {peak_reserved_mb:.2f} MB")
        print(_gpu_process_table(processes, os.getpid()), end="")
        print("--------------------------------------------------\n")
        self.last_memory_reserved_mb = current_reserved_mb

    def _register_tmp_dir(self, d: str):
        if d and os.path.isdir(d):
            self._tmp_dirs.add(d); print(f"[DEBUG] Registrado tmp dir: {d}")

    def _register_tmp_file(self, f: str):
        if f and os.path.exists(f):
            self._tmp_files.add(f); print(f"[DEBUG] Registrado tmp file: {f}")

    def finalize(self, keep_paths=None, extra_paths=None, clear_gpu=True):
        print("[DEBUG] Finalize: iniciando limpeza...")
        keep = set(keep_paths or []); extras = set(extra_paths or [])
        removed_files = 0
        for f in list(self._tmp_files | extras):
            try:
                if f not in keep and os.path.isfile(f):
                    os.remove(f); removed_files += 1; print(f"[DEBUG] Removido arquivo tmp: {f}")
            except Exception as e:
                print(f"[DEBUG] Falha removendo arquivo {f}: {e}")
            finally:
                self._tmp_files.discard(f)
        removed_dirs = 0
        for d in list(self._tmp_dirs):
            try:
                if d not in keep and os.path.isdir(d):
                    shutil.rmtree(d, ignore_errors=True); removed_dirs += 1; print(f"[DEBUG] Removido diretório tmp: {d}")
            except Exception as e:
                print(f"[DEBUG] Falha removendo diretório {d}: {e}")
            finally:
                self._tmp_dirs.discard(d)
        print(f"[DEBUG] Finalize: arquivos removidos={removed_files}, dirs removidos={removed_dirs}")
        gc.collect()
        try:
            if clear_gpu and torch.cuda.is_available():
                torch.cuda.empty_cache()
                try:
                    torch.cuda.ipc_collect()
                except Exception:
                    pass
        except Exception as e:
            print(f"[DEBUG] Finalize: limpeza GPU falhou: {e}")
        try:
            self._log_gpu_memory("Após finalize")
        except Exception as e:
            print(f"[DEBUG] Log GPU pós-finalize falhou: {e}")

    def _load_config(self):
        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 in candidates:
            if cfg.exists():
                print(f"[DEBUG] Config selecionada: {cfg}")
                with open(cfg, "r") as file:
                    return yaml.safe_load(file)
        cfg = base / "ltxv-13b-0.9.8-distilled-fp8.yaml"
        print(f"[DEBUG] Config fallback: {cfg}")
        with open(cfg, "r") as file:
            return yaml.safe_load(file)

    def _load_models(self):
        """
        Carrega os modelos de forma inteligente:
        1. Tenta resolver o caminho do cache local (rápido, sem rede).
        2. Se o arquivo não for encontrado localmente, baixa como fallback.
        Garante que o serviço possa iniciar mesmo que o setup.py não tenha sido executado.
        """
        t0 = time.perf_counter()
        LTX_REPO = "Lightricks/LTX-Video"

        print("[DEBUG] Resolvendo caminhos dos modelos de forma inteligente...")

        # --- Função Auxiliar para Carregamento Inteligente ---
        def get_or_download_model(repo_id, filename, description):
            try:
                # hf_hub_download é a ferramenta certa aqui. Ela verifica o cache PRIMEIRO.
                # Se o arquivo estiver no cache, retorna o caminho instantaneamente (após uma verificação rápida de metadados).
                # Se não estiver no cache, ela o baixa.
                print(f"[DEBUG] Verificando {description}: {filename}...")
                model_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    # Forçar o uso de um cache específico se necessário
                    cache_dir=os.getenv("HF_HOME_CACHE"),
                    token=os.getenv("HF_TOKEN")
                )
                print(f"[DEBUG] Caminho do {description} resolvido com sucesso.")
                return model_path
            except Exception as e:
                print("\n" + "="*80)
                print(f"[ERRO CRÍTICO] Falha ao obter o modelo '{filename}'.")
                print(f"Detalhe do erro: {e}")
                print("Verifique sua conexão com a internet ou o estado do cache do Hugging Face.")
                print("="*80 + "\n")
                sys.exit(1)

        # --- Checkpoint Principal ---
        checkpoint_filename = self.config["checkpoint_path"]
        distilled_model_path = get_or_download_model(
            LTX_REPO, checkpoint_filename, "checkpoint principal"
        )
        self.config["checkpoint_path"] = distilled_model_path
        
        # --- Upscaler Espacial ---
        upscaler_filename = self.config["spatial_upscaler_model_path"]
        spatial_upscaler_path = get_or_download_model(
            LTX_REPO, upscaler_filename, "upscaler espacial"
        )
        self.config["spatial_upscaler_model_path"] = spatial_upscaler_path

        # --- Construção dos Pipelines ---
        print("\n[DEBUG] Construindo pipeline a partir dos caminhos resolvidos...")
        pipeline = create_ltx_video_pipeline(
            ckpt_path=self.config["checkpoint_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",
            enhance_prompt=False,
            prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"],
            prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"],
        )
        print("[DEBUG] Pipeline pronto.")

        latent_upsampler = None
        if self.config.get("spatial_upscaler_model_path"):
            print("[DEBUG] Construindo latent_upsampler...")
            latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
            print("[DEBUG] Upsampler pronto.")
            
        print(f"[DEBUG] _load_models() tempo total={time.perf_counter()-t0:.3f}s")
        return pipeline, latent_upsampler

    
    def _load_models_old(self):
        t0 = time.perf_counter()
        LTX_REPO = "Lightricks/LTX-Video"
        print("[DEBUG] Baixando checkpoint principal...")
        distilled_model_path = hf_hub_download(
            repo_id=LTX_REPO,
            filename=self.config["checkpoint_path"],
            local_dir=os.getenv("HF_HOME"),
            cache_dir=os.getenv("HF_HOME_CACHE"),
            token=os.getenv("HF_TOKEN"),
        )
        self.config["checkpoint_path"] = distilled_model_path
        print(f"[DEBUG] Checkpoint em: {distilled_model_path}")

        print("[DEBUG] Baixando upscaler espacial...")
        spatial_upscaler_path = hf_hub_download(
            repo_id=LTX_REPO,
            filename=self.config["spatial_upscaler_model_path"],
            local_dir=os.getenv("HF_HOME"),
            cache_dir=os.getenv("HF_HOME_CACHE"),
            token=os.getenv("HF_TOKEN")
        )
        self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
        print(f"[DEBUG] Upscaler em: {spatial_upscaler_path}")

        print("[DEBUG] Construindo pipeline...")
        pipeline = create_ltx_video_pipeline(
            ckpt_path=self.config["checkpoint_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",
            enhance_prompt=False,
            prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"],
            prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"],
        )
        print("[DEBUG] Pipeline pronto.")

        latent_upsampler = None
        if self.config.get("spatial_upscaler_model_path"):
            print("[DEBUG] Construindo latent_upsampler...")
            latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
            print("[DEBUG] Upsampler pronto.")
        print(f"[DEBUG] _load_models() tempo total={time.perf_counter()-t0:.3f}s")
        return pipeline, latent_upsampler

    def _promote_fp8_weights_to_bf16(self, module):
        if not isinstance(module, torch.nn.Module):
            print("[DEBUG] Promoção FP8→BF16 ignorada: alvo não é nn.Module.")
            return
        f8 = getattr(torch, "float8_e4m3fn", None)
        if f8 is None:
            print("[DEBUG] torch.float8_e4m3fn indisponível.")
            return
        p_cnt = b_cnt = 0
        for _, p in module.named_parameters(recurse=True):
            try:
                if p.dtype == f8:
                    with torch.no_grad():
                        p.data = p.data.to(torch.bfloat16); p_cnt += 1
            except Exception:
                pass
        for _, b in module.named_buffers(recurse=True):
            try:
                if hasattr(b, "dtype") and b.dtype == f8:
                    b.data = b.data.to(torch.bfloat16); b_cnt += 1
            except Exception:
                pass
        print(f"[DEBUG] FP8→BF16: params_promoted={p_cnt}, buffers_promoted={b_cnt}")

    @torch.no_grad()
    def _upsample_latents_internal(self, latents: torch.Tensor) -> torch.Tensor:
        """
        Lógica extraída diretamente da LTXMultiScalePipeline para upscale de latentes.
        """
        if not self.latent_upsampler:
            raise ValueError("Latent Upsampler não está carregado.")

        # Garante que os modelos estejam no dispositivo correto
        self.latent_upsampler.to(self.device)
        self.pipeline.vae.to(self.device)
        print(f"[DEBUG-UPSAMPLE] Shape de entrada: {tuple(latents.shape)}")
        latents = un_normalize_latents(latents, self.pipeline.vae, vae_per_channel_normalize=True)
        upsampled_latents = self.latent_upsampler(latents)
        upsampled_latents = normalize_latents(upsampled_latents, self.pipeline.vae, vae_per_channel_normalize=True)
        print(f"[DEBUG-UPSAMPLE] Shape de saída: {tuple(upsampled_latents.shape)}")
        
        return upsampled_latents
        

    def _apply_precision_policy(self):
        prec = str(self.config.get("precision", "")).lower()
        self.runtime_autocast_dtype = torch.float32
        print(f"[DEBUG] Aplicando política de precisão: {prec}")
        if prec == "float8_e4m3fn":
            self.runtime_autocast_dtype = torch.bfloat16
            force_promote = os.getenv("LTXV_FORCE_BF16_ON_FP8", "0") == "1"
            print(f"[DEBUG] FP8 detectado. force_promote={force_promote}")
            if force_promote and hasattr(torch, "float8_e4m3fn"):
                try:
                    self._promote_fp8_weights_to_bf16(self.pipeline)
                except Exception as e:
                    print(f"[DEBUG] Promoção FP8→BF16 na pipeline falhou: {e}")
                try:
                    if self.latent_upsampler:
                        self._promote_fp8_weights_to_bf16(self.latent_upsampler)
                except Exception as e:
                    print(f"[DEBUG] Promoção FP8→BF16 no upsampler falhou: {e}")
        elif prec == "bfloat16":
            self.runtime_autocast_dtype = torch.bfloat16
        elif prec == "mixed_precision":
            self.runtime_autocast_dtype = torch.float16
        else:
            self.runtime_autocast_dtype = torch.float32

    def _prepare_conditioning_tensor(self, filepath, height, width, padding_values):
        print(f"[DEBUG] Carregando condicionamento: {filepath}")
        tensor = load_image_to_tensor_with_resize_and_crop(filepath, height, width)
        tensor = torch.nn.functional.pad(tensor, padding_values)
        out = tensor.to(self.device, dtype=self.runtime_autocast_dtype) if self.device == "cuda" else tensor.to(self.device)
        print(f"[DEBUG] Cond shape={tuple(out.shape)} dtype={out.dtype} device={out.device}")
        return out


    def _dividir_latentes_por_tamanho(self, latents_brutos, num_latente_por_chunk: int, overlap: int = 1):
        """
        Divide o tensor de latentes em chunks com tamanho definido em número de latentes.
        
        Args:
            latents_brutos: tensor [B, C, T, H, W]
            num_latente_por_chunk: número de latentes por chunk
            overlap: número de frames que se sobrepõem entre chunks
        
        Returns:
            List[tensor]: lista de chunks cloneados
        """
        sum_latent = latents_brutos.shape[2]
        chunks = []
    
        if num_latente_por_chunk >= sum_latent:
            return [latents_brutos]
    
        n_chunks = (sum_latent) // num_latente_por_chunk 
        steps = sum_latent//n_chunks
        print("==========PODA CAUSAL[start:stop-1]==========")
        print(f"[DEBUG] TOTAL LATENTES = {sum_latent}")
        print(f"[DEBUG] LATENTES min por chunk = {num_latente_por_chunk}")
        print(f"[DEBUG] Número de chunks = {n_chunks}")
        if n_chunks > 1:
            i=0
            while i < n_chunks:
                if i>0:
                    dow=0
                else:
                    dow=0
                start = (num_latente_por_chunk*i)
                end = (start+num_latente_por_chunk+(overlap+1))
                if i+1 < n_chunks:
                    chunk = latents_brutos[:, :, start-(dow):end, :, :].clone().detach()
                    print(f"[DEBUG] chunk{i+1}[:, :, {start-dow}:{end}, :, :] = {chunk.shape[2]}")
                else:
                    chunk = latents_brutos[:, :, start-(dow):, :, :].clone().detach()
                    print(f"[DEBUG] chunk{i+1}[:, :, {start-(dow)}:, :, :] = {chunk.shape[2]}")
                chunks.append(chunk)
                i+=1
        else:
            print(f"[DEBUG] numero chunks minimo ")
            print(f"[DEBUG] latents_brutos[:, :, :, :, :] = {latents_brutos.shape[2]}")
            chunks.append(latents_brutos)
        print("\n\n================PODA CAUSAL=================")
        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",
            video_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return int(result.stdout.strip())    

    def _gerar_lista_com_transicoes(self, pasta: str, video_paths: list[str], crossfade_frames: int = 8) -> list[str]:
        """
        Gera uma nova lista de vídeos aplicando transições suaves (blend frame a frame)
        seguindo exatamente a lógica linear de Carlos.
        """
        import os, subprocess, shutil
    
        poda = crossfade_frames
        total_partes = len(video_paths)
        video_fade_fim = None
        video_fade_ini = None
        nova_lista = []

        if crossfade_frames == 0:
            print("\n\n[DEBUG] CROSSFADE_FRAMES=0 Ship concatenation causal")
            return video_paths
        

        print("\n\n===========CONCATECAO CAUSAL=============")
     
        print(f"[DEBUG] Iniciando pipeline com {total_partes} vídeos e {poda} frames de crossfade")

        
        for i in range(total_partes):
            base = video_paths[i]
        
            # --- PODA ---
            video_podado = os.path.join(pasta, f"{base}_podado_{i}.mp4")
            

            if i<total_partes-1:
                end_frame = self._get_total_frames(base) - poda
            else:
                end_frame = self._get_total_frames(base)
            
            if i>0:
                start_frame = poda
            else:
                start_frame = 0
                
            cmd_fim = (
               f'ffmpeg -y -hide_banner -loglevel error -i "{base}" '
               f'-vf "trim=start_frame={start_frame}:end_frame={end_frame},setpts=PTS-STARTPTS" '
               f'-an "{video_podado}"'
            )
            subprocess.run(cmd_fim, shell=True, check=True)
            
        
            # --- FADE_INI --- 
            if i > 0:
                video_fade_ini = os.path.join(pasta, f"{base}_fade_ini_{i}.mp4")
                cmd_ini = (
                    f'ffmpeg -y -hide_banner -loglevel error -i "{base}" '
                    f'-vf "trim=end_frame={poda},setpts=PTS-STARTPTS" -an "{video_fade_ini}"'
                )
                subprocess.run(cmd_ini, shell=True, check=True)
                
            # --- TRANSIÇÃO ---
            if video_fade_fim and video_fade_ini:
                video_fade  = os.path.join(pasta, f"transicao_{i}_{i+1}.mp4")
                cmd_blend = (
                    f'ffmpeg -y -hide_banner -loglevel error '
                    f'-i "{video_fade_fim}" -i "{video_fade_ini}" '
                    f'-filter_complex "[0:v][1:v]blend=all_expr=\'A*(1-T/{poda})+B*(T/{poda})\',format=yuv420p" '
                    f'-frames:v {poda} "{video_fade}"'
                )
                subprocess.run(cmd_blend, shell=True, check=True)
                print(f"[DEBUG] transicao adicionada {i}/{i+1} {self._get_total_frames(video_fade)} frames ✅")
                nova_lista.append(video_fade)
                
            # --- FADE_FIM ---
            if i<=total_partes-1:
                video_fade_fim = os.path.join(pasta, f"{base}_fade_fim_{i}.mp4")
                cmd_fim = (
                    f'ffmpeg -y -hide_banner -loglevel error -i "{base}" '
                    f'-vf "trim=start_frame={end_frame-poda},setpts=PTS-STARTPTS" -an "{video_fade_fim}"'
                )
                subprocess.run(cmd_fim, shell=True, check=True)
           
            nova_lista.append(video_podado)
            print(f"[DEBUG] Video podado {i+1} adicionado {self._get_total_frames(video_podado)} frames ✅")



        print("===========CONCATECAO CAUSAL=============")
        print(f"[DEBUG] {nova_lista}")
        return nova_lista
        
    def _concat_mp4s_no_reencode(self, mp4_list: List[str], out_path: str):
        """
        Concatena múltiplos MP4s sem reencode usando o demuxer do ffmpeg.
        ATENÇÃO: todos os arquivos precisam ter mesmo codec, fps, resolução etc.
        """
        if not mp4_list or len(mp4_list) < 2:
            raise ValueError("Forneça pelo menos dois arquivos MP4 para concatenar.")

        
        # Cria lista temporária para o ffmpeg
        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".txt") 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}"
        print(f"[DEBUG] Concat: {cmd}")
    
        try:
            subprocess.check_call(shlex.split(cmd))
        finally:
            try:
                os.remove(list_path)
            except Exception:
                pass   
   

    # ==============================================================================
    # --- FUNÇÃO GENERATE COMPLETA E ATUALIZADA ---
    # ==============================================================================
    def generate(
        self,
        prompt,
        negative_prompt,
        mode="text-to-video",
        start_image_filepath=None,
        middle_image_filepath=None,
        middle_frame_number=None,
        middle_image_weight=1.0,
        end_image_filepath=None,
        end_image_weight=1.0,
        input_video_filepath=None,
        height=512,
        width=704,
        duration=2.0,
        frames_to_use=9,
        seed=42,
        randomize_seed=True,
        guidance_scale=3.0,
        improve_texture=True,
        progress_callback=None,
        external_decode=True,
    ):
        t_all = time.perf_counter()
        print(f"[DEBUG] generate() begin mode={mode} external_decode={external_decode} improve_texture={improve_texture}")
        if self.device == "cuda":
            torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
        self._log_gpu_memory("Início da Geração")

        # --- Setup Inicial (como antes) ---
        if mode == "image-to-video" and not start_image_filepath:
            raise ValueError("A imagem de início é obrigatória para o modo image-to-video")
        used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
        seed_everething(used_seed); print(f"[DEBUG] Seed usado: {used_seed}")
        FPS = 24.0; MAX_NUM_FRAMES = 2570
        target_frames_rounded = round(duration * FPS)
        n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
        actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 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)
        
        conditioning_items = []
        if mode == "image-to-video":
            start_tensor = self._prepare_conditioning_tensor(start_image_filepath, height, width, padding_values)
            conditioning_items.append(ConditioningItem(start_tensor, 0, 1.0))
            if middle_image_filepath and middle_frame_number is not None:
                middle_tensor = self._prepare_conditioning_tensor(middle_image_filepath, height, width, padding_values)
                safe_middle_frame = max(0, min(int(middle_frame_number), actual_num_frames - 1))
                conditioning_items.append(ConditioningItem(middle_tensor, safe_middle_frame, float(middle_image_weight)))
            if end_image_filepath:
                end_tensor = self._prepare_conditioning_tensor(end_image_filepath, height, width, padding_values)
                last_frame_index = actual_num_frames - 1
                conditioning_items.append(ConditioningItem(end_tensor, last_frame_index, float(end_image_weight)))
            print(f"[DEBUG] Conditioning items: {len(conditioning_items)}")

        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 if conditioning_items else None,
            "media_items": None,
            "decode_timestep": self.config["decode_timestep"],
            "decode_noise_scale": self.config["decode_noise_scale"],
            "stochastic_sampling": self.config["stochastic_sampling"],
            "image_cond_noise_scale": 0.01,
            "is_video": True,
            "vae_per_channel_normalize": True,
            "mixed_precision": (self.config["precision"] == "mixed_precision"),
            "offload_to_cpu": False,
            "enhance_prompt": False,
            "skip_layer_strategy": SkipLayerStrategy.AttentionValues,
        }
        
        latents = None
        latents_list = []
        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
            

        try:
            ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
            with ctx:
                if improve_texture:
                    if not self.latent_upsampler:
                        raise ValueError("Upscaler espacial não carregado, mas 'improve_texture' está ativo.")

                    first_pass_kwargs = call_kwargs.copy()
    
                    # --- ETAPA 1: GERAÇÃO BASE (FIRST PASS) ---
                    print("\n--- INICIANDO ETAPA 1: GERAÇÃO BASE (FIRST PASS) ---")
                    t_pass1 = time.perf_counter()

                    first_pass_config = self.config.get("first_pass", {}).copy()
                    downscale_factor = self.config.get("downscale_factor", 0.6666666)
                    vae_scale_factor = self.pipeline.vae_scale_factor # Geralmente 8
                    
                    x_width = int(width_padded * downscale_factor)
                    downscaled_width = x_width - (x_width % vae_scale_factor)
                    x_height = int(height_padded * downscale_factor)
                    downscaled_height = x_height - (x_height % vae_scale_factor)
                    print(f"[DEBUG] First Pass Dims: Original Pad ({width_padded}x{height_padded}) -> Downscaled ({downscaled_width}x{downscaled_height})")

                    first_pass_kwargs.update({
                        **first_pass_config
                    })
                    
                    first_pass_kwargs.update({
                        "output_type": "latent",
                        "width": downscaled_width,
                        "height": downscaled_height,
                        "guidance_scale": float(guidance_scale),
                    })
                    
                    print(f"[DEBUG] First Pass: Gerando em {downscaled_width}x{downscaled_height}...")
                    base_latents = self.pipeline(**first_pass_kwargs).images
                    log_tensor_info(base_latents, "Latentes Base (First Pass)")
                    print(f"[DEBUG] First Pass concluída em {time.perf_counter() - t_pass1:.2f}s")
                    
                    # --- ETAPA 2: UPSCALE DOS LATENTES ---
                    print("\n--- INICIANDO ETAPA 2: UPSCALE DOS LATENTES ---")
                    t_upscale = time.perf_counter()
                    
                    upsampled_latents = self._upsample_latents_internal(base_latents)
                    upsampled_latents = adain_filter_latent(latents=upsampled_latents, reference_latents=base_latents)
                    log_tensor_info(upsampled_latents, "Latentes Pós-Upscale")
                    print(f"[DEBUG] Upscale de Latentes concluído em {time.perf_counter() - t_upscale:.2f}s")
                    del base_latents; gc.collect(); torch.cuda.empty_cache()

                    par = 0
                    latents_cpu_up = upsampled_latents.detach().to("cpu", non_blocking=True)
                    torch.cuda.empty_cache()
                    try:
                         torch.cuda.ipc_collect()
                    except Exception:
                         pass
                
                    latents_parts_up = self._dividir_latentes_por_tamanho(latents_cpu_up,40,0)

                    print("\n\n--- INICIANDO ETAPA 3: REFINAMENTO DE TEXTURA (SECOND PASS) ---")

                    cc = 1
                    for latents in latents_parts_up:

                        t_pass2 = time.perf_counter()
                        print("\n\n#########################################")
                        # # --- ETAPA 3: REFINAMENTO DE TEXTURA (SECOND PASS) ---
                        print(f"\n--- INICIANDO ETAPA 3/{cc} ")

                        second_pass_kwargs = first_pass_config.copy()
                        second_pass_config = self.config.get("second_pass", {}).copy()

                        second_pass_width = downscaled_width * 2
                        second_pass_height = downscaled_height * 2
                        print(f"[DEBUG] Second Pass Dims: Target ({second_pass_width}x{second_pass_height})")
                        
                        num_latent_frames_part = latents.shape[2]
                        log_tensor_info(latents, "Latentes input (Pre-Pós-Second Pass)")
                    
                        vae_temporal_scale = self.pipeline.video_scale_factor # Geralmente 4 ou 8
                        num_pixel_frames_part = ((num_latent_frames_part - 1) * vae_temporal_scale) + 1
                        print(f"[DEBUG] Parte: {num_latent_frames_part - 1} latentes -> {num_pixel_frames_part} frames de pixel (alvo)")

                        second_pass_kwargs.update({
                           **second_pass_config
                        })
                        
                        second_pass_kwargs.update({
                           "output_type": "latent",
                           "width": second_pass_width,
                           "height": second_pass_height,
                           "num_frames": num_pixel_frames_part,
                           "latents": latents, # O tensor upscaled
                           "guidance_scale": float(guidance_scale),
                        })
                    
                        print(f"[DEBUG] Second Pass: Refinando em {width_padded}x{height_padded}...")
                        final_latents = self.pipeline(**second_pass_kwargs).images
                        log_tensor_info(final_latents, "Latentes Finais (Pós-Second Pass)")
                        print(f"[DEBUG] Second part Pass concluída em {time.perf_counter() - t_pass2:.2f}s")
                        
                        latents_list.append(final_latents)
                        cc+=1
                        print("#########################################")
                      
                    print("\n\n--- FIM ETAPA 3: REFINAMENTO DE TEXTURA (SECOND PASS) ---")

                else: # Geração de etapa única
                    print("\n--- INICIANDO GERAÇÃO DE ETAPA ÚNICA ---")
                    t_single = time.perf_counter()
                    single_pass_kwargs = call_kwargs.copy()
                    single_pass_kwargs.update(self.config.get("first_pass", {}))
                    single_pass_kwargs["guidance_scale"] = float(guidance_scale)
                    single_pass_kwargs["output_type"] = "latent"
                    
                    # Remove keys that might conflict or are not used in single pass / handled by above
                    single_pass_kwargs.pop("num_inference_steps", None) 
                    single_pass_kwargs.pop("first_pass", None) 
                    single_pass_kwargs.pop("second_pass", None)
                    single_pass_kwargs.pop("downscale_factor", None)
                    
                    latents = self.pipeline(**single_pass_kwargs).images
                    log_tensor_info(latents, "Latentes Finais (Etapa Única)")
                    print(f"[DEBUG] Etapa única concluída em {time.perf_counter() - t_single:.2f}s")

                    latents_list.append(latents)
            
            # --- ETAPA FINAL: DECODIFICAÇÃO E CODIFICAÇÃO MP4 ---
            print("\n--- INICIANDO ETAPA FINAL: DECODIFICAÇÃO E MONTAGEM ---")
            
            temp_dir = tempfile.mkdtemp(prefix="ltxv_"); self._register_tmp_dir(temp_dir)
            results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
    
            
            partes_mp4 = []
            par = 0
            
            for latents_vae in latents_list:
            
                latents_cpu_vae = latents_vae.detach().to("cpu", non_blocking=True)
                torch.cuda.empty_cache()
                try:
                    torch.cuda.ipc_collect()
                except Exception:
                    pass
                    
                latents_parts_vae = self._dividir_latentes_por_tamanho(latents_cpu_vae,4,1)
                
                for latents in latents_parts_vae:
                    print(f"[DEBUG] Partição {par}: {tuple(latents.shape)}")
           
                    par = par + 1
                    output_video_path = os.path.join(temp_dir, f"output_{used_seed}_{par}.mp4")
                    final_output_path = None
        
                    print("[DEBUG] Decodificando bloco de latentes com VAE → tensor de pixels...")
                    # Usar manager com timestep por item; previne target_shape e rota NoneType.decode
                    pixel_tensor = vae_manager_singleton.decode(
                        latents.to(self.device, non_blocking=True),
                        decode_timestep=float(self.config.get("decode_timestep", 0.05))
                    )
                    log_tensor_info(pixel_tensor, "Pixel tensor (VAE saída)")
        
                    print("[DEBUG] Codificando MP4 a partir do tensor de pixels (bloco inteiro)...")
                    video_encode_tool_singleton.save_video_from_tensor(
                        pixel_tensor,
                        output_video_path,
                        fps=call_kwargs["frame_rate"],
                        progress_callback=progress_callback
                    )
        
                    candidate = os.path.join(results_dir, f"output_par_{par}.mp4")
                    try:
                        shutil.move(output_video_path, candidate)
                        final_output_path = candidate
                        print(f"[DEBUG] MP4 parte {par} movido para {final_output_path}")
                        partes_mp4.append(final_output_path)
                        
                    except Exception as e:
                        final_output_path = output_video_path
                        print(f"[DEBUG] Falha no move; usando tmp como final: {e}")
    
            
            total_partes = len(partes_mp4)
            if (total_partes>1):
                final_vid = os.path.join(results_dir, f"concat_fim_{used_seed}.mp4")
                partes_mp4_fade = self._gerar_lista_com_transicoes(pasta=results_dir, video_paths=partes_mp4, crossfade_frames=0)
                self._concat_mp4s_no_reencode(partes_mp4_fade, final_vid)
            else:
                final_vid = partes_mp4[0]
            
            
            self._log_gpu_memory("Fim da Geração")
            return final_vid, used_seed
            

        except Exception as e:
            print("[DEBUG] EXCEÇÃO NA GERAÇÃO:")
            print("".join(traceback.format_exception(type(e), e, e.__traceback__)))
            raise
        finally:
            try:
                del latents
            except Exception:
                pass
            try:
                del multi_scale_pipeline
            except Exception:
                pass

            gc.collect()
            try:
                if self.device == "cuda":
                    torch.cuda.empty_cache()
                    try:
                        torch.cuda.ipc_collect()
                    except Exception:
                        pass
            except Exception as e:
                print(f"[DEBUG] Limpeza GPU no finally falhou: {e}")

            try:
                self.finalize(keep_paths=[])
            except Exception as e:
                print(f"[DEBUG] finalize() no finally falhou: {e}")

print("Criando instância do VideoService. O carregamento do modelo começará agora...")
video_generation_service = VideoService()