api/ltx_video_complete.py

# ==============================================================================
# ltx_video_service_with_gpu_pools.py
# VideoService com Multi-GPU Pool Manager Integrado
# ==============================================================================
# Arquitetura:
# - GPU 0 e 1: Pipeline + Upscaler (geração/refinamento de latentes)
# - GPU 2 e 3: VAE Decode (decodificação de latentes para pixels)
# ==============================================================================

import os
import sys
import gc
import yaml
import time
import json
import random
import shutil
import warnings
import tempfile
import traceback
import subprocess
import threading
import queue
from pathlib import Path
from typing import List, Dict, Optional, Tuple, Union
from dataclasses import dataclass
from enum import Enum
import cv2
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from einops import rearrange
from huggingface_hub import hf_hub_download
from safetensors import safe_open

# --- Configurações ---
ENABLE_MEMORY_OPTIMIZATION = os.getenv("ADUC_MEMORY_OPTIMIZATION", "1").lower() in ["1", "true", "yes"]
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
from huggingface_hub import logging as hf_logging
hf_logging.set_verbosity_error()

# --- Importações de managers ---
from managers.vae_manager import vae_manager_singleton
from tools.video_encode_tool import video_encode_tool_singleton

# --- Constantes Globais ---
LTXV_DEBUG = True
LTXV_FRAME_LOG_EVERY = 8
DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
RESULTS_DIR = Path("/app/output")
DEFAULT_FPS = 24.0

# ==============================================================================
# SETUP E IMPORTAÇÕES DO REPOSITÓRIO
# ==============================================================================

def _run_setup_script():
    """Executa o script setup.py se o repositório LTX-Video não existir."""
    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

    print(f"[DEBUG] Repositório não encontrado em {LTX_VIDEO_REPO_DIR}. Executando setup.py...")
    try:
        subprocess.run([sys.executable, setup_script_path], check=True, capture_output=True, text=True)
        print("[DEBUG] Script 'setup.py' concluído com sucesso.")
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] Falha ao executar 'setup.py' (código {e.returncode}).\nOutput:\n{e.stdout}\n{e.stderr}")
        sys.exit(1)

def add_deps_to_path(repo_path: Path):
    """Adiciona o diretório do repositório ao sys.path para importações locais."""
    resolved_path = str(repo_path.resolve())
    if resolved_path not in sys.path:
        sys.path.insert(0, resolved_path)
        if LTXV_DEBUG:
            print(f"[DEBUG] Adicionado ao sys.path: {resolved_path}")

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

# --- Importações Dependentes do Path Adicionado ---
from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
from ltx_video.pipelines.pipeline_ltx_video import adain_filter_latent
from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXVideoPipeline
from transformers import T5EncoderModel, T5Tokenizer, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
from ltx_video.models.transformers.transformer3d import Transformer3DModel
from ltx_video.schedulers.rf import RectifiedFlowScheduler
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
import ltx_video.pipelines.crf_compressor as crf_compressor

# ==============================================================================
# GPU POOL MANAGER - Sistema Multi-GPU
# ==============================================================================

class GPUPoolType(Enum):
    """Tipos de pools de GPU disponíveis"""
    GENERATION = "generation"  # Pipeline + Upscaler
    DECODE = "decode"          # VAE Decode


@dataclass
class GPUTask:
    """Representa uma tarefa a ser executada em uma GPU"""
    task_id: str
    task_fn: callable
    args: tuple
    kwargs: dict
    result_queue: queue.Queue


@dataclass
class GPUWorker:
    """Representa um worker de GPU individual"""
    worker_id: int
    device_id: str
    pool_type: GPUPoolType
    thread: Optional[threading.Thread] = None
    is_busy: bool = False


class GPUPoolManager:
    """
    Gerenciador de pools de GPU para distribuição de tarefas.
    
    Arquitetura:
    - Pool 1 (GENERATION): 2 GPUs para pipeline + upscaler
    - Pool 2 (DECODE): 2 GPUs para VAE decode
    """
    
    def __init__(
        self,
        generation_devices: List[str] = None,
        decode_devices: List[str] = None,
        max_queue_size: int = 10
    ):
        """Inicializa o gerenciador de pools."""
        self.generation_devices = generation_devices or ["cuda:0", "cuda:1"]
        self.decode_devices = decode_devices or ["cuda:2", "cuda:3"]
        
        self.generation_queue = queue.Queue(maxsize=max_queue_size)
        self.decode_queue = queue.Queue(maxsize=max_queue_size)
        
        self.generation_workers: List[GPUWorker] = []
        self.decode_workers: List[GPUWorker] = []
        
        self._shutdown = False
        self._lock = threading.Lock()
        
        self.stats = {
            "generation_tasks_completed": 0,
            "decode_tasks_completed": 0,
            "generation_tasks_failed": 0,
            "decode_tasks_failed": 0,
        }
        
        self._initialize_workers()
    
    def _initialize_workers(self):
        """Inicializa todos os workers de GPU"""
        print("[GPU Pool Manager] Inicializando workers...")
        
        for i, device in enumerate(self.generation_devices):
            worker = GPUWorker(
                worker_id=i,
                device_id=device,
                pool_type=GPUPoolType.GENERATION
            )
            worker.thread = threading.Thread(
                target=self._worker_loop,
                args=(worker, self.generation_queue),
                daemon=True
            )
            worker.thread.start()
            self.generation_workers.append(worker)
            print(f"  ✓ Generation Worker {i} iniciado em {device}")
        
        for i, device in enumerate(self.decode_devices):
            worker = GPUWorker(
                worker_id=i,
                device_id=device,
                pool_type=GPUPoolType.DECODE
            )
            worker.thread = threading.Thread(
                target=self._worker_loop,
                args=(worker, self.decode_queue),
                daemon=True
            )
            worker.thread.start()
            self.decode_workers.append(worker)
            print(f"  ✓ Decode Worker {i} iniciado em {device}")
        
        print(f"[GPU Pool Manager] {len(self.generation_workers)} workers de GERAÇÃO e {len(self.decode_workers)} workers de DECODE ativos.\n")
    
    def _worker_loop(self, worker: GPUWorker, task_queue: queue.Queue):
        """Loop principal de um worker."""
        print(f"[Worker {worker.worker_id}:{worker.device_id}] Aguardando tarefas ({worker.pool_type.value})...")
        
        while not self._shutdown:
            try:
                task: GPUTask = task_queue.get(timeout=1.0)
                
                with self._lock:
                    worker.is_busy = True
                
                print(f"[Worker {worker.worker_id}:{worker.device_id}] Executando tarefa '{task.task_id}'...")
                
                try:
                    torch.cuda.set_device(worker.device_id)
                    result = task.task_fn(
                        worker.device_id,
                        *task.args,
                        **task.kwargs
                    )
                    task.result_queue.put(("success", result))
                    
                    with self._lock:
                        if worker.pool_type == GPUPoolType.GENERATION:
                            self.stats["generation_tasks_completed"] += 1
                        else:
                            self.stats["decode_tasks_completed"] += 1
                    
                    print(f"[Worker {worker.worker_id}:{worker.device_id}] Tarefa '{task.task_id}' concluída com sucesso.")
                
                except Exception as e:
                    print(f"[Worker {worker.worker_id}:{worker.device_id}] ERRO na tarefa '{task.task_id}': {e}")
                    import traceback
                    traceback.print_exc()
                    
                    task.result_queue.put(("error", str(e)))
                    
                    with self._lock:
                        if worker.pool_type == GPUPoolType.GENERATION:
                            self.stats["generation_tasks_failed"] += 1
                        else:
                            self.stats["decode_tasks_failed"] += 1
                
                finally:
                    with self._lock:
                        worker.is_busy = False
                    task_queue.task_done()
                    torch.cuda.empty_cache()
            
            except queue.Empty:
                continue
    
    def submit_generation_task(
        self,
        task_id: str,
        task_fn: callable,
        *args,
        **kwargs
    ) -> queue.Queue:
        """Submete uma tarefa de GERAÇÃO ao pool."""
        result_queue = queue.Queue(maxsize=1)
        task = GPUTask(
            task_id=task_id,
            task_fn=task_fn,
            args=args,
            kwargs=kwargs,
            result_queue=result_queue
        )
        
        print(f"[GPU Pool Manager] Submetendo tarefa de GERAÇÃO: '{task_id}'")
        self.generation_queue.put(task)
        return result_queue
    
    def submit_decode_task(
        self,
        task_id: str,
        task_fn: callable,
        *args,
        **kwargs
    ) -> queue.Queue:
        """Submete uma tarefa de DECODE ao pool."""
        result_queue = queue.Queue(maxsize=1)
        task = GPUTask(
            task_id=task_id,
            task_fn=task_fn,
            args=args,
            kwargs=kwargs,
            result_queue=result_queue
        )
        
        print(f"[GPU Pool Manager] Submetendo tarefa de DECODE: '{task_id}'")
        self.decode_queue.put(task)
        return result_queue
    
    def get_result(self, result_queue: queue.Queue, timeout: Optional[float] = None):
        """Aguarda e retorna o resultado de uma tarefa."""
        status, result = result_queue.get(timeout=timeout)
        
        if status == "error":
            raise Exception(f"Tarefa falhou: {result}")
        
        return result
    
    def submit_and_wait_generation(
        self,
        task_id: str,
        task_fn: callable,
        *args,
        timeout: Optional[float] = None,
        **kwargs
    ):
        """Submete uma tarefa de geração e aguarda o resultado (bloqueante)."""
        result_queue = self.submit_generation_task(task_id, task_fn, *args, **kwargs)
        return self.get_result(result_queue, timeout=timeout)
    
    def submit_and_wait_decode(
        self,
        task_id: str,
        task_fn: callable,
        *args,
        timeout: Optional[float] = None,
        **kwargs
    ):
        """Submete uma tarefa de decode e aguarda o resultado (bloqueante)."""
        result_queue = self.submit_decode_task(task_id, task_fn, *args, **kwargs)
        return self.get_result(result_queue, timeout=timeout)
    
    def wait_all(self):
        """Aguarda todas as tarefas pendentes serem concluídas"""
        print("[GPU Pool Manager] Aguardando conclusão de todas as tarefas...")
        self.generation_queue.join()
        self.decode_queue.join()
        print("[GPU Pool Manager] Todas as tarefas concluídas.")
    
    def get_stats(self) -> dict:
        """Retorna estatísticas de uso do pool"""
        with self._lock:
            return {
                **self.stats,
                "generation_queue_size": self.generation_queue.qsize(),
                "decode_queue_size": self.decode_queue.qsize(),
                "generation_workers_busy": sum(1 for w in self.generation_workers if w.is_busy),
                "decode_workers_busy": sum(1 for w in self.decode_workers if w.is_busy),
            }
    
    def print_stats(self):
        """Imprime estatísticas formatadas"""
        stats = self.get_stats()
        print("\n" + "="*60)
        print("GPU POOL MANAGER - ESTATÍSTICAS")
        print("="*60)
        print(f"Generation Pool:")
        print(f"  - Tarefas Concluídas: {stats['generation_tasks_completed']}")
        print(f"  - Tarefas Falhadas:   {stats['generation_tasks_failed']}")
        print(f"  - Workers Ocupados:   {stats['generation_workers_busy']}/{len(self.generation_workers)}")
        print(f"  - Fila:               {stats['generation_queue_size']} tarefas")
        print(f"\nDecode Pool:")
        print(f"  - Tarefas Concluídas: {stats['decode_tasks_completed']}")
        print(f"  - Tarefas Falhadas:   {stats['decode_tasks_failed']}")
        print(f"  - Workers Ocupados:   {stats['decode_workers_busy']}/{len(self.decode_workers)}")
        print(f"  - Fila:               {stats['decode_queue_size']} tarefas")
        print("="*60 + "\n")
    
    def shutdown(self):
        """Encerra todos os workers"""
        print("[GPU Pool Manager] Encerrando...")
        self._shutdown = True
        
        for worker in self.generation_workers + self.decode_workers:
            if worker.thread:
                worker.thread.join(timeout=5.0)
        
        print("[GPU Pool Manager] Encerrado.")


# Singleton global
_gpu_pool_manager_instance: Optional[GPUPoolManager] = None


def get_gpu_pool_manager(
    generation_devices: List[str] = None,
    decode_devices: List[str] = None,
    force_reinit: bool = False
) -> GPUPoolManager:
    """Retorna a instância singleton do GPUPoolManager."""
    global _gpu_pool_manager_instance
    
    if _gpu_pool_manager_instance is None or force_reinit:
        if _gpu_pool_manager_instance and force_reinit:
            _gpu_pool_manager_instance.shutdown()
        
        _gpu_pool_manager_instance = GPUPoolManager(
            generation_devices=generation_devices,
            decode_devices=decode_devices
        )
    
    return _gpu_pool_manager_instance

# ==============================================================================
# FUNÇÕES AUXILIARES DE PROCESSAMENTO
# ==============================================================================

def debug_log(message: str):
    """Log condicional baseado em LTXV_DEBUG"""
    if LTXV_DEBUG:
        print(f"[DEBUG] {message}")

def load_image_cv2(image_path: str, target_height: int, target_width: int) -> np.ndarray:
    """Carrega uma imagem usando OpenCV e redimensiona"""
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"Não foi possível carregar a imagem: {image_path}")
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
    return image

def normalize_image(image: np.ndarray) -> np.ndarray:
    """Normaliza imagem para [-1, 1]"""
    image = image.astype(np.float32) / 127.5 - 1.0
    return image

def denormalize_image(image: np.ndarray) -> np.ndarray:
    """Desnormaliza imagem de [-1, 1] para [0, 255]"""
    image = (image + 1.0) * 127.5
    return np.clip(image, 0, 255).astype(np.uint8)

# ==============================================================================
# CLASSE PRINCIPAL DO SERVIÇO DE VÍDEO COM GPU POOLS
# ==============================================================================

class VideoService:
    """
    Serviço de Geração de Vídeos com LTX Video e Multi-GPU Pool Manager.
    
    Arquitetura de GPUs:
    - GPU 0 e 1: Pipeline + Upscaler (GENERATION pool)
    - GPU 2 e 3: VAE Decode (DECODE pool)
    """
    
    def __init__(self):
        """Inicializa o serviço com GPU Pools"""
        print("[VideoService] Inicializando com Multi-GPU Pools...")
        
        # Inicializa o pool manager
        self.gpu_pool = get_gpu_pool_manager(
            generation_devices=["cuda:0", "cuda:1"],
            decode_devices=["cuda:2", "cuda:3"]
        )
        
        # Carrega configuração
        self.config = self._load_config("ltxv-13b-0.9.8-distilled-fp8.yaml")
        
        # Carrega modelos (template que será clonado para cada GPU)
        self.pipeline_template, self.latent_upsampler_template = self._load_models_from_hub()
        
        # Inicializa pipelines em cada GPU de geração
        self.generation_models = {}
        for device in ["cuda:0", "cuda:1"]:
            self.generation_models[device] = self._clone_pipeline_to_device(device)
        
        # Inicializa VAE em cada GPU de decode
        self.decode_models = {}
        for device in ["cuda:2", "cuda:3"]:
            self.decode_models[device] = self._clone_vae_to_device(device)
        
        # Configurações de tempo de execução
        self.runtime_autocast_dtype = self._get_precision_dtype()
        
        # Anexa pipeline ao vae_manager_singleton
        vae_manager_singleton.attach_pipeline(
            self.pipeline_template,
            device="cuda:0",
            autocast_dtype=self.runtime_autocast_dtype
        )
        
        # Rastreamento de seed
        self.used_seed = None
        self.tmp_dir = None
        self._register_tmp_dir()
        
        print("[VideoService] Inicializado com sucesso!")
        print("[VideoService] Pools de GPU ativos:")
        print("[VideoService]   - Geração: cuda:0, cuda:1")
        print("[VideoService]   - Decode: cuda:2, cuda:3")
    
    def _clone_pipeline_to_device(self, device: str) -> Dict:
        """Clona a pipeline para um dispositivo específico"""
        print(f"  Clonando pipeline para {device}...")
        pipeline = {
            'transformer': self.pipeline_template.transformer.to(device),
            'text_encoder': self.pipeline_template.text_encoder.to(device),
            'scheduler': self.pipeline_template.scheduler,
            'tokenizer': self.pipeline_template.tokenizer,
            'patchifier': self.pipeline_template.patchifier,
        }
        
        if self.latent_upsampler_template:
            pipeline['upsampler'] = self.latent_upsampler_template.to(device)
        
        return pipeline
    
    def _clone_vae_to_device(self, device: str) -> torch.nn.Module:
        """Clona o VAE para um dispositivo específico"""
        print(f"  Clonando VAE para {device}...")
        vae = self.pipeline_template.vae.to(device)
        vae.eval()
        return vae
    
    # ==============================================================================
    # FUNÇÕES WORKER PARA POOL MANAGER
    # ==============================================================================
    
    def _generate_latents_worker(
        self,
        device_id: str,
        prompt: str,
        negative_prompt: str,
        height: int,
        width: int,
        num_frames: int,
        guidance_scale: float,
        seed: int,
        conditioning_items: Optional[List] = None
    ) -> torch.Tensor:
        """Worker para geração de latentes (roda em cuda:0 ou cuda:1)"""
        print(f"    [Generation Worker] Gerando latentes em {device_id}")
        
        generator = torch.Generator(device=device_id).manual_seed(seed)
        
        with torch.autocast(device_type='cuda', dtype=self.runtime_autocast_dtype):
            kwargs = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "height": height,
                "width": width,
                "num_frames": num_frames,
                "frame_rate": int(DEFAULT_FPS),
                "generator": generator,
                "output_type": "latent",
                "guidance_scale": float(guidance_scale),
                "conditioning_items": conditioning_items,
                **self.config.get("first_pass", {})
            }
            
            latents = self.pipeline_template(**kwargs).images
        
        # Aplica upsampler se disponível
        if 'upsampler' in self.generation_models[device_id]:
            latents = self._upsample_and_filter_latents(
                latents,
                self.generation_models[device_id]['upsampler'],
                device_id
            )
        
        return latents.cpu()
    
    def _refine_latents_worker(
        self,
        device_id: str,
        latents: torch.Tensor,
        prompt: str,
        negative_prompt: str,
        guidance_scale: float,
        seed: int,
        conditioning_items: Optional[List] = None
    ) -> torch.Tensor:
        """Worker para refinamento de latentes (roda em cuda:0 ou cuda:1)"""
        print(f"    [Refine Worker] Refinando latentes em {device_id}")
        
        latents = latents.to(device_id)
        
        with torch.autocast(device_type='cuda', dtype=self.runtime_autocast_dtype):
            refine_height = latents.shape[3] * 8  # vae_scale_factor
            refine_width = latents.shape[4] * 8
            
            kwargs = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "height": refine_height,
                "width": refine_width,
                "frame_rate": int(DEFAULT_FPS),
                "num_frames": latents.shape[2],
                "latents": latents,
                "guidance_scale": float(guidance_scale),
                "output_type": "latent",
                "generator": torch.Generator(device=device_id).manual_seed(seed),
                "conditioning_items": conditioning_items,
                **self.config.get("second_pass", {})
            }
            
            refined_latents = self.pipeline_template(**kwargs).images
        
        return refined_latents.cpu()
    
    def _decode_latents_worker(
        self,
        device_id: str,
        latents: torch.Tensor,
        decode_timestep: float = 0.05
    ) -> torch.Tensor:
        """Worker para decodificação de latentes (roda em cuda:2 ou cuda:3)"""
        print(f"    [Decode Worker] Decodificando em {device_id} (shape: {latents.shape})")
        
        latents = latents.to(device_id)
        vae = self.decode_models[device_id]
        
        with torch.no_grad():
            with torch.autocast(device_type='cuda', dtype=self.runtime_autocast_dtype):
                pixel_tensor = vae_manager_singleton.decode(
                    latents,
                    decode_timestep=decode_timestep
                )
        
        return pixel_tensor.cpu()
    
    # ==============================================================================
    # MÉTODOS DE PREPARAÇÃO DE DADOS
    # ==============================================================================
    
    def _load_image_to_tensor_with_resize_and_crop(
        self,
        image_path: str,
        target_height: int,
        target_width: int,
        padding_values: tuple = (0, 0, 0)
    ) -> torch.Tensor:
        """Carrega uma imagem, redimensiona e converte para tensor"""
        image = load_image_cv2(image_path, target_height, target_width)
        image = normalize_image(image)
        tensor = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
        return tensor
    
    def _prepare_conditioning_tensor(
        self,
        image_path: str,
        target_height: int,
        target_width: int,
        padding_values: tuple = (0, 0, 0)
    ) -> torch.Tensor:
        """Prepara tensor de condicionamento de uma imagem"""
        return self._load_image_to_tensor_with_resize_and_crop(
            image_path,
            target_height,
            target_width,
            padding_values
        )
    
    def _prepare_conditioning_tensor_from_path(self, image_path: str) -> torch.Tensor:
        """Prepara tensor de condicionamento com resolução padrão"""
        return self._prepare_conditioning_tensor(image_path, 512, 768, (0, 0, 0))
    
    # ==============================================================================
    # MÉTODOS DE CÁLCULO E PROCESSAMENTO
    # ==============================================================================
    
    def _calculate_downscaled_dims(self, height: int, width: int) -> Tuple[int, int]:
        """Calcula dimensões reduzidas para primeira passagem"""
        downscale_factor = 4
        return height // downscale_factor, width // downscale_factor
    
    def _calculate_dynamic_cuts(
        self,
        total_latents: int,
        min_chunk_size: int = 8,
        overlap: int = 2
    ) -> Tuple[List[Tuple[int, int]], List[int]]:
        """Calcula pontos de corte dinâmicos para chunks com overlap"""
        cut_points = []
        segment_sizes = []
        
        start = 0
        while start < total_latents:
            end = min(start + min_chunk_size, total_latents)
            cut_points.append((start, end))
            segment_sizes.append(end - start)
            
            if end >= total_latents:
                break
            
            start = end - overlap
        
        return cut_points, segment_sizes
    
    def _split_latents_with_overlap(
        self,
        latents: torch.Tensor,
        chunk_size: int = 8,
        overlap: int = 2
    ) -> List[torch.Tensor]:
        """Divide latentes em chunks com overlap"""
        chunks = []
        start = 0
        total_frames = latents.shape[2]
        
        while start < total_frames:
            end = min(start + chunk_size, total_frames)
            chunk = latents[:, :, start:end, :, :]
            chunks.append(chunk)
            
            if end >= total_frames:
                break
            
            start = end - overlap
        
        return chunks
    
    def _merge_chunks_with_overlap(
        self,
        chunks: List[torch.Tensor],
        overlap: int = 2
    ) -> torch.Tensor:
        """Costura chunks removendo overlap"""
        if len(chunks) == 1:
            return chunks[0]
        
        overlap_pixels = overlap * 8  # 8 = VAE scale factor
        
        result_parts = [chunks[0][:, :, :-overlap_pixels, :, :]]
        
        for chunk in chunks[1:-1]:
            result_parts.append(chunk[:, :, overlap_pixels:-overlap_pixels, :, :])
        
        if len(chunks) > 1:
            result_parts.append(chunks[-1][:, :, overlap_pixels:, :, :])
        
        return torch.cat(result_parts, dim=2)
    
    def _stitch_dynamic_chunks(
        self,
        pixel_chunks: List[torch.Tensor],
        segment_sizes: List[int],
        macro_overlap: int = 2
    ) -> torch.Tensor:
        """Costura chunks dinâmicos com tratamento de overlap"""
        if len(pixel_chunks) == 1:
            return pixel_chunks[0]
        
        overlap_frames = macro_overlap * 8
        stitched_parts = []
        
        for i, chunk in enumerate(pixel_chunks):
            if i == 0:
                stitched_parts.append(chunk[:, :, :-overlap_frames, :, :])
            elif i == len(pixel_chunks) - 1:
                stitched_parts.append(chunk[:, :, overlap_frames:, :, :])
            else:
                stitched_parts.append(chunk[:, :, overlap_frames:-overlap_frames, :, :])
        
        return torch.cat(stitched_parts, dim=2)
    
    def _upsample_and_filter_latents(
        self,
        latents: torch.Tensor,
        upsampler: torch.nn.Module,
        device: str
    ) -> torch.Tensor:
        """Aplica upsampler e filtro aos latentes"""
        latents = latents.to(device)
        
        with torch.no_grad():
            with torch.autocast(device_type='cuda', dtype=self.runtime_autocast_dtype):
                upsampled = upsampler(latents)
                filtered = adain_filter_latent(upsampled, latents)
        
        return filtered
    
    # ==============================================================================
    # MÉTODOS DE GERAÇÃO E REFINAMENTO (USANDO POOL MANAGER)
    # ==============================================================================
    
    def generate_low_resolution(
        self,
        prompt: str,
        negative_prompt: str,
        height: int,
        width: int,
        duration_secs: float,
        guidance_scale: float,
        seed: Optional[int] = None,
        image_filepaths: Optional[List[str]] = None
    ) -> Tuple[str, int]:
        """Gera vídeo em baixa resolução usando pool de geração"""
        print("[INFO] Iniciando geração em baixa resolução (modo paralelo)...")
        
        used_seed = seed or random.randint(0, 2**32 - 1)
        self._seed_everething(used_seed)
        
        actual_num_frames = int(round(duration_secs * DEFAULT_FPS))
        downscaled_height, downscaled_width = self._calculate_downscaled_dims(height, width)
        
        conditioning_items = []
        if image_filepaths:
            for filepath in image_filepaths:
                cond_tensor = self._prepare_conditioning_tensor(
                    filepath,
                    downscaled_height,
                    downscaled_width,
                    (0, 0, 0)
                )
                conditioning_items.append(ConditioningItem(cond_tensor, 0, 1.0))
        
        # Submete tarefa de geração ao pool
        task_id = f"gen_lowres_{used_seed}"
        latents = self.gpu_pool.submit_and_wait_generation(
            task_id=task_id,
            task_fn=self._generate_latents_worker,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=downscaled_height,
            width=downscaled_width,
            num_frames=(actual_num_frames // 8) + 1,
            guidance_scale=guidance_scale,
            seed=used_seed,
            conditioning_items=conditioning_items if conditioning_items else None,
            timeout=600
        )
        
        tensor_path = self._save_latents_to_disk(latents, "latents_low_res", used_seed)
        
        print("[SUCCESS] Geração de baixa resolução concluída!")
        self.used_seed = used_seed
        
        return tensor_path, used_seed
    
    def refine_texture_only(
        self,
        latents_path: str,
        prompt: str,
        negative_prompt: str,
        guidance_scale: float,
        seed: int,
        image_filepaths: Optional[List[str]] = None,
        macro_chunk_size: int = 8,
        macro_overlap: int = 2
    ) -> Tuple[str, str, torch.Tensor]:
        """Refina e decodifica latentes usando ambos os pools em paralelo"""
        print("[INFO] Iniciando refinamento e decodificação paralela...")
        
        initial_latents = torch.load(latents_path).cpu()
        total_latents = initial_latents.shape[2]
        height = initial_latents.shape[3] * 8
        width = initial_latents.shape[4] * 8
        
        cut_points, segment_sizes = self._calculate_dynamic_cuts(
            total_latents,
            min_chunk_size=macro_chunk_size,
            overlap=macro_overlap
        )
        
        print(f"  Processando {len(cut_points)} chunks em paralelo...")
        
        # Prepara conditioning se fornecido
        conditioning_items = []
        if image_filepaths:
            for filepath in image_filepaths:
                cond_tensor = self._prepare_conditioning_tensor(
                    filepath,
                    height,
                    width,
                    (0, 0, 0)
                )
                conditioning_items.append(ConditioningItem(cond_tensor, 0, 1.0))
        
        pixel_results = []
        
        for i, (start, end) in enumerate(cut_points):
            chunk_id = f"chunk_{i}_seed_{seed}"
            latent_chunk = initial_latents[:, :, start:end, :, :]
            
            # ETAPA 1: Refinar latentes (pool de geração)
            print(f"\n  [{i+1}/{len(cut_points)}] Refinando chunk {start}-{end}...")
            refined_latents = self.gpu_pool.submit_and_wait_generation(
                task_id=f"refine_{chunk_id}",
                task_fn=self._refine_latents_worker,
                latents=latent_chunk,
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=guidance_scale,
                seed=seed + i,
                conditioning_items=conditioning_items if conditioning_items else None,
                timeout=600
            )
            
            # ETAPA 2: Decodificar latentes (pool de decode)
            print(f"  [{i+1}/{len(cut_points)}] Decodificando chunk {start}-{end}...")
            pixel_tensor = self.gpu_pool.submit_and_wait_decode(
                task_id=f"decode_{chunk_id}",
                task_fn=self._decode_latents_worker,
                latents=refined_latents,
                decode_timestep=float(self.config.get("decode_timestep", 0.05)),
                timeout=300
            )
            
            pixel_results.append(pixel_tensor)
            
            del refined_latents
            torch.cuda.empty_cache()
        
        # Costura resultados
        print("\n  Costurando chunks finais...")
        final_pixel_tensor = self._stitch_dynamic_chunks(
            pixel_results,
            segment_sizes,
            macro_overlap
        )
        
        final_video_path = self._save_video_from_tensor(
            final_pixel_tensor,
            "final_video",
            seed
        )
        
        print(f"[SUCCESS] Vídeo final salvo em: {final_video_path}")
        self.gpu_pool.print_stats()
        
        return final_video_path, latents_path, final_pixel_tensor
    
    def apply_secondary_refinement(
        self,
        initial_latents_path: str,
        prompt: str,
        negative_prompt: str,
        guidance_scale: float,
        seed: int,
        image_filepaths: Optional[List[str]] = None
    ) -> str:
        """Aplica refinamento secundário em múltiplos chunks"""
        print("[INFO] Aplicando refinamento secundário...")
        
        initial_latents = torch.load(initial_latents_path).cpu()
        total_latents = initial_latents.shape[2]
        
        # Divide em chunks maiores
        macro_chunk_size = 16
        macro_overlap = 2
        
        cut_points, segment_sizes = self._calculate_dynamic_cuts(
            total_latents,
            min_chunk_size=macro_chunk_size,
            overlap=macro_overlap
        )
        
        height = initial_latents.shape[3] * 8
        width = initial_latents.shape[4] * 8
        
        conditioning_items = []
        if image_filepaths:
            for filepath in image_filepaths:
                cond_tensor = self._prepare_conditioning_tensor(
                    filepath, height, width, (0, 0, 0)
                )
                conditioning_items.append(ConditioningItem(cond_tensor, 0, 1.0))
        
        print(f"  Refinando {len(cut_points)} chunks...")
        
        # Submete TODAS as tarefas de refinamento
        refine_queues = []
        for i, (start, end) in enumerate(cut_points):
            latent_chunk = initial_latents[:, :, start:end, :, :]
            
            queue = self.gpu_pool.submit_generation_task(
                task_id=f"refine_macro_{i}",
                task_fn=self._refine_latents_worker,
                latents=latent_chunk,
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=guidance_scale,
                seed=seed + i,
                conditioning_items=conditioning_items if conditioning_items else None
            )
            refine_queues.append((i, queue))
        
        # Processa decodes conforme refinamentos ficam prontos
        print(f"\n  Decodificando chunks refinados...")
        decode_queues = []
        
        for i, refine_queue in refine_queues:
            refined_latents = self.gpu_pool.get_result(refine_queue, timeout=600)
            print(f"    ✓ Chunk {i} refinado")
            
            decode_queue = self.gpu_pool.submit_decode_task(
                task_id=f"decode_macro_{i}",
                task_fn=self._decode_latents_worker,
                latents=refined_latents,
                decode_timestep=float(self.config.get("decode_timestep", 0.05))
            )
            decode_queues.append((i, decode_queue))
        
        # Aguarda todos os decodes
        print(f"\n  Aguardando conclusão de todos os decodes...")
        pixel_results = []
        
        for i, decode_queue in decode_queues:
            pixel_tensor = self.gpu_pool.get_result(decode_queue, timeout=300)
            pixel_results.append(pixel_tensor)
            print(f"    ✓ Chunk {i} decodificado")
        
        # Costura resultados finais
        print(f"\n  Costurando resultado final...")
        final_pixel_tensor = self._stitch_dynamic_chunks(
            pixel_results,
            segment_sizes,
            macro_overlap
        )
        
        final_video_path = self._save_video_from_tensor(
            final_pixel_tensor,
            "refined_final_video",
            seed
        )
        
        print(f"[SUCCESS] Vídeo refinado salvo em: {final_video_path}")
        self.gpu_pool.print_stats()
        
        return final_video_path
    
    def encode_latents_to_mp4(
        self,
        pixel_tensor: torch.Tensor,
        output_path: str,
        fps: float = 24.0
    ) -> str:
        """Codifica tensor de pixels em arquivo MP4"""
        print(f"[INFO] Codificando vídeo para MP4: {output_path}")
        
        # Desnormaliza
        pixel_tensor = (pixel_tensor + 1.0) / 2.0 * 255.0
        pixel_tensor = torch.clamp(pixel_tensor, 0, 255)
        
        # Converte para formato de vídeo
        video_encode_tool_singleton.encode_video_from_tensor(
            pixel_tensor,
            output_path,
            fps=fps
        )
        
        print(f"[SUCCESS] Vídeo codificado: {output_path}")
        return output_path
    
    # ==============================================================================
    # MÉTODOS DE CONFIGURAÇÃO E CARREGAMENTO
    # ==============================================================================
    
    def _load_config(self, config_file: str) -> Dict:
        """Carrega configuração YAML"""
        config_path = LTX_VIDEO_REPO_DIR / "configs" / config_file
        
        if not config_path.exists():
            print(f"[WARNING] Arquivo de config não encontrado: {config_path}")
            return {}
        
        with open(config_path, "r") as f:
            config = yaml.safe_load(f)
        
        return config or {}
    
    def _load_models_from_hub(self) -> Tuple[LTXVideoPipeline, Optional[LatentUpsampler]]:
        """Carrega modelos do Hugging Face Hub"""
        print("[INFO] Carregando modelos do Hub...")
        
        # Carrega pipeline
        pipeline = LTXVideoPipeline.from_pretrained(
            "Lightricks/LTX-Video",
            torch_dtype=torch.bfloat16
        )
        
        # Carrega upsampler (opcional)
        try:
            upsampler = LatentUpsampler.from_pretrained(
                "Lightricks/LTX-Video",
                torch_dtype=torch.bfloat16
            )
        except Exception as e:
            print(f"[WARNING] Upsampler não disponível: {e}")
            upsampler = None
        
        print("[SUCCESS] Modelos carregados com sucesso!")
        return pipeline, upsampler
    
    def _move_models_to_device(self):
        """Move modelos para dispositivo principal (não usado com pools)"""
        # Implementado no _clone_pipeline_to_device
        pass
    
    def _get_precision_dtype(self) -> torch.dtype:
        """Retorna tipo de dados de precisão baseado em disponibilidade"""
        if torch.cuda.is_available():
            device_props = torch.cuda.get_device_properties(0)
            if device_props.major >= 8:  # A100, H100, etc.
                return torch.bfloat16
        
        return torch.float16
    
    # ==============================================================================
    # MÉTODOS AUXILIARES DE SALVAMENTO E GERENCIAMENTO
    # ==============================================================================
    
    def _save_latents_to_disk(
        self,
        latents: torch.Tensor,
        prefix: str,
        seed: int
    ) -> str:
        """Salva latentes em arquivo .pt"""
        filename = f"{prefix}_{seed}.pt"
        filepath = self.tmp_dir / filename
        
        torch.save(latents, filepath)
        print(f"  Latentes salvos: {filepath}")
        
        return str(filepath)
    
    def _save_video_from_tensor(
        self,
        pixel_tensor: torch.Tensor,
        prefix: str,
        seed: int
    ) -> str:
        """Salva tensor de pixels como vídeo MP4"""
        filename = f"{prefix}_{seed}.mp4"
        filepath = RESULTS_DIR / filename
        
        RESULTS_DIR.mkdir(parents=True, exist_ok=True)
        
        self.encode_latents_to_mp4(pixel_tensor, str(filepath), fps=DEFAULT_FPS)
        
        print(f"  Vídeo salvo: {filepath}")
        return str(filepath)
    
    def _finalize(self):
        """Finaliza o serviço e libera recursos"""
        print("[INFO] Finalizando VideoService...")
        
        self.gpu_pool.print_stats()
        self.gpu_pool.shutdown()
        
        if self.tmp_dir and self.tmp_dir.exists():
            shutil.rmtree(self.tmp_dir)
            print(f"  Diretório temporário removido: {self.tmp_dir}")
        
        # Limpa memória CUDA
        torch.cuda.empty_cache()
        gc.collect()
        
        print("[SUCCESS] VideoService finalizado!")
    
    def _seed_everething(self, seed: int):
        """Define seed para reproducibilidade"""
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    
    def _register_tmp_dir(self):
        """Registra diretório temporário para salvamento de latentes"""
        self.tmp_dir = Path(tempfile.mkdtemp(prefix="ltx_video_"))
        print(f"  Diretório temporário: {self.tmp_dir}")

# ==============================================================================
# PONTO DE ENTRADA E EXEMPLO DE USO
# ==============================================================================

if __name__ == "__main__":
    print("\n" + "="*80)
    print("LTX VIDEO SERVICE - Multi-GPU Pool Manager")
    print("="*80 + "\n")
    
    try:
        # Inicializa o serviço
        print("Criando instância do VideoService...")
        video_service = VideoService()
        
        # Exemplo 1: Geração de baixa resolução
        print("\n[EXEMPLO 1] Geração de baixa resolução...")
        latents_path, seed = video_service.generate_low_resolution(
            prompt="A beautiful sunset over the ocean",
            negative_prompt="blurry, low quality",
            height=512,
            width=768,
            duration_secs=2.0,
            guidance_scale=3.0,
            seed=42,
            image_filepaths=None
        )
        
        # Exemplo 2: Refinamento e decodificação
        print("\n[EXEMPLO 2] Refinamento e decodificação...")
        video_path, latents_path, final_tensor = video_service.refine_texture_only(
            latents_path=latents_path,
            prompt="A beautiful sunset over the ocean",
            negative_prompt="blurry, low quality",
            guidance_scale=3.0,
            seed=seed,
            image_filepaths=None,
            macro_chunk_size=8,
            macro_overlap=2
        )
        
        print(f"\n✓ Vídeo final gerado: {video_path}")
        
    except KeyboardInterrupt:
        print("\n\n[INFO] Interrompido pelo usuário.")
    except Exception as e:
        print(f"\n\n[ERROR] Erro na execução: {e}")
        import traceback
        traceback.print_exc()
    finally:
        if 'video_service' in locals():
            video_service._finalize()
    
    print("\n" + "="*80)
    print("Execução concluída")
    print("="*80 + "\n")