Upload 2 files

api/ltx_server_refactored_complete (1).py

@@ -0,0 +1,288 @@
+# FILE: api/ltx_server_refactored_complete.py
+# DESCRIPTION: Final orchestrator for LTX-Video generation.
+# This version internalizes conditioning item preparation, accepting a raw
+# list of media items directly in its main generation function for maximum simplicity and encapsulation.
+
+import gc
+import json
+import logging
+import os
+import shutil
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import yaml
+import numpy as np
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+# ==============================================================================
+# --- SETUP E IMPORTAÇÕES DO PROJETO ---
+# ==============================================================================
+
+# Configuração de logging e supressão de warnings
+import warnings
+warnings.filterwarnings("ignore")
+logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+log_level = os.environ.get("ADUC_LOG_LEVEL", "INFO").upper()
+logging.basicConfig(level=log_level, format='[%(levelname)s] [%(name)s] %(message)s')
+
+# --- Constantes de Configuração ---
+DEPS_DIR = Path("/data")
+LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
+RESULTS_DIR = Path("/app/output")
+DEFAULT_FPS = 24.0
+FRAMES_ALIGNMENT = 8
+LTX_REPO_ID = "Lightricks/LTX-Video"
+
+
+# --- Módulos da nossa Arquitetura ---
+try:
+    from api.gpu_manager import gpu_manager
+    from api.vae_server import vae_server_singleton
+    from tools.video_encode_tool import video_encode_tool_singleton
+    from api.ltx.ltx_utils import build_ltx_pipeline_on_cpu, seed_everything
+    from api.ltx_pool_manager import LatentConditioningItem
+    from api.utils.debug_utils import log_function_io
+except ImportError as e:
+    logging.critical(f"A crucial import from the local API/architecture failed. Error: {e}", exc_info=True)
+    sys.exit(1)
+
+# ==============================================================================
+# --- CLASSE DE SERVIÇO (O ORQUESTRADOR) ---
+# ==============================================================================
+
+class VideoService:
+    """
+    Orchestrates the high-level logic of video generation, with internalized
+    conditioning item preparation.
+    """
+
+    @log_function_io
+    def __init__(self):
+        t0 = time.time()
+        logging.info("Initializing VideoService Orchestrator...")
+        RESULTS_DIR.mkdir(parents=True, exist_ok=True)
+
+        target_main_device_str = str(gpu_manager.get_ltx_device())
+        target_vae_device_str = str(gpu_manager.get_ltx_vae_device())
+        logging.info(f"LTX allocated to devices: Main='{target_main_device_str}', VAE='{target_vae_device_str}'")
+
+        self.config = self._load_config()
+        self._resolve_model_paths_from_cache()
+
+        self.pipeline, self.latent_upsampler = build_ltx_pipeline_on_cpu(self.config)
+
+        self.main_device = torch.device("cpu")
+        self.vae_device = torch.device("cpu")
+        self.move_to_device(main_device_str=target_main_device_str, vae_device_str=target_vae_device_str)
+
+        self._apply_precision_policy()
+        logging.info(f"VideoService ready. Startup time: {time.time() - t0:.2f}s")
+
+    def _load_config(self) -> Dict:
+        """Loads the YAML configuration file."""
+        config_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
+        with open(config_path, "r") as file:
+            return yaml.safe_load(file)
+
+    def _resolve_model_paths_from_cache(self):
+        """Finds the absolute paths to model files in the cache and updates the in-memory config."""
+        logging.info("Resolving model paths from Hugging Face cache...")
+        cache_dir = os.environ.get("HF_HOME")
+        try:
+            main_ckpt_path = hf_hub_download(repo_id=LTX_REPO_ID, filename=self.config["checkpoint_path"], cache_dir=cache_dir)
+            self.config["checkpoint_path"] = main_ckpt_path
+            if self.config.get("spatial_upscaler_model_path"):
+                upscaler_path = hf_hub_download(repo_id=LTX_REPO_ID, filename=self.config["spatial_upscaler_model_path"], cache_dir=cache_dir)
+                self.config["spatial_upscaler_model_path"] = upscaler_path
+        except Exception as e:
+            logging.critical(f"Failed to resolve model paths. Ensure setup.py ran correctly. Error: {e}", exc_info=True)
+            sys.exit(1)
+
+    @log_function_io
+    def move_to_device(self, main_device_str: str, vae_device_str: str):
+        """Moves pipeline components to their designated target devices."""
+        target_main_device = torch.device(main_device_str)
+        target_vae_device = torch.device(vae_device_str)
+        self.main_device = target_main_device
+        self.vae_device = target_vae_device
+        self.pipeline.to(self.main_device)
+        self.pipeline.vae.to(self.vae_device)
+        if self.latent_upsampler: self.latent_upsampler.to(self.main_device)
+        logging.info("LTX models successfully moved to target devices.")
+
+    def move_to_cpu(self):
+        """Moves all LTX components to CPU to free VRAM for other services."""
+        self.move_to_device(main_device_str="cpu", vae_device_str="cpu")
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+
+    def finalize(self):
+        """Cleans up GPU memory after a generation task."""
+        gc.collect()
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        try: torch.cuda.ipc_collect();
+        except Exception: pass
+
+    # ==========================================================================
+    # --- LÓGICA DE NEGÓCIO: ORQUESTRADOR PÚBLICO UNIFICADO ---
+    # ==========================================================================
+
+    @log_function_io
+    def generate_low_resolution(
+        self,
+        prompt_list: List[str],
+        initial_media_items: Optional[List[Tuple[Union[str, Image.Image, torch.Tensor], int, float]]] = None,
+        **kwargs
+    ) -> Tuple[Optional[str], Optional[str], Optional[int]]:
+        """
+        [UNIFIED ORCHESTRATOR] Generates a low-resolution video from a prompt and a raw list of media items.
+        """
+        logging.info("Starting unified low-resolution generation...")
+        used_seed = self._get_random_seed()
+        seed_everything(used_seed)
+        logging.info(f"Using randomly generated seed: {used_seed}")
+
+        if not prompt_list: raise ValueError("Prompt is empty or contains no valid lines.")
+        
+        is_narrative = len(prompt_list) > 1
+        num_chunks = len(prompt_list)
+        total_frames = self._calculate_aligned_frames(kwargs.get("duration", 4.0))
+        frames_per_chunk = max(FRAMES_ALIGNMENT, (total_frames // num_chunks // FRAMES_ALIGNMENT) * FRAMES_ALIGNMENT)
+        overlap_frames = 9 if is_narrative else 0
+        
+        initial_conditions = []
+        if initial_media_items:
+            logging.info("Preparing initial conditioning items from raw media list...")
+            initial_conditions = vae_server_singleton.generate_conditioning_items(
+                media_items=[item[0] for item in initial_media_items],
+                target_frames=[item[1] for item in initial_media_items],
+                strengths=[item[2] for item in initial_media_items],
+                target_resolution=(kwargs['height'], kwargs['width'])
+            )
+        
+        temp_latent_paths = []
+        overlap_condition_item: Optional[LatentConditioningItem] = None
+        
+        try:
+            for i, chunk_prompt in enumerate(prompt_list):
+                logging.info(f"Processing scene {i+1}/{num_chunks}: '{chunk_prompt[:50]}...'")
+                
+                if i < num_chunks - 1:
+                    current_frames_base = frames_per_chunk
+                else:
+                    processed_frames_base = (num_chunks - 1) * frames_per_chunk
+                    current_frames_base = total_frames - processed_frames_base
+                
+                current_frames = current_frames_base + (overlap_frames if i > 0 else 0)
+                current_frames = self._align(current_frames, alignment_rule='n*8+1')
+
+                current_conditions = initial_conditions if i == 0 else []
+                if overlap_condition_item: current_conditions.append(overlap_condition_item)
+
+                chunk_latents = self._generate_single_chunk_low(
+                    prompt=chunk_prompt, num_frames=current_frames, seed=used_seed + i,
+                    conditioning_items=current_conditions, **kwargs
+                )
+                if chunk_latents is None: raise RuntimeError(f"Failed to generate latents for scene {i+1}.")
+
+                if is_narrative and i < num_chunks - 1:
+                    overlap_latents = chunk_latents[:, :, -overlap_frames:, :, :].clone()
+                    overlap_condition_item = LatentConditioningItem(
+                        latent_tensor=overlap_latents.cpu(),
+                        media_frame_number=0,
+                        conditioning_strength=1.0
+                    )
+                
+                if i > 0: chunk_latents = chunk_latents[:, :, overlap_frames:, :, :]
+                
+                chunk_path = RESULTS_DIR / f"temp_chunk_{i}_{used_seed}.pt"
+                torch.save(chunk_latents.cpu(), chunk_path)
+                temp_latent_paths.append(chunk_path)
+            
+            base_filename = "narrative_video" if is_narrative else "single_video"
+            all_tensors_cpu = [torch.load(p) for p in temp_latent_paths]
+            final_latents = torch.cat(all_tensors_cpu, dim=2)
+
+            video_path, latents_path = self._finalize_generation(final_latents, base_filename, used_seed)
+            return video_path, latents_path, used_seed
+        except Exception as e:
+            logging.error(f"Error during unified generation: {e}", exc_info=True)
+            return None, None, None
+        finally:
+            for path in temp_latent_paths:
+                if path.exists(): path.unlink()
+            self.finalize()
+
+    # ==========================================================================
+    # --- UNIDADES DE TRABALHO E HELPERS INTERNOS ---
+    # ==========================================================================
+    
+    def _log_conditioning_items(self, items: List[Union[ConditioningItem, LatentConditioningItem]]):
+        """Logs detailed information about a list of ConditioningItem objects."""
+        if logging.getLogger().isEnabledFor(logging.DEBUG):
+            # (Lógica de logging para debug)
+            pass
+
+    @log_function_io
+    def _generate_single_chunk_low(self, **kwargs) -> Optional[torch.Tensor]:
+        """[WORKER] Calls the pipeline to generate a single chunk of latents."""
+        # (A lógica desta função permanece a mesma)
+        pass # Placeholder
+
+    @log_function_io
+    def _finalize_generation(self, final_latents: torch.Tensor, base_filename: str, seed: int) -> Tuple[str, str]:
+        """Consolidates latents, decodes them to video, and saves final artifacts."""
+        logging.info("Finalizing generation: decoding latents to video.")
+        final_latents_path = RESULTS_DIR / f"latents_{base_filename}_{seed}.pt"
+        torch.save(final_latents, final_latents_path)
+        logging.info(f"Final latents saved to: {final_latents_path}")
+        
+        pixel_tensor = vae_server_singleton.decode_to_pixels(
+            final_latents, decode_timestep=float(self.config.get("decode_timestep", 0.05))
+        )
+        video_path = self._save_and_log_video(pixel_tensor, f"{base_filename}_{seed}")
+        return str(video_path), str(final_latents_path)
+    
+    def _apply_ui_overrides(self, config_dict: Dict, overrides: Dict):
+        """Applies advanced settings from the UI to a config dictionary."""
+        # (Lógica de overrides da UI permanece a mesma)
+        pass # Placeholder
+
+    def _save_and_log_video(self, pixel_tensor: torch.Tensor, base_filename: str) -> Path:
+        """Saves a pixel tensor (on CPU) to an MP4 file."""
+        # (Lógica de salvar vídeo permanece a mesma)
+        pass # Placeholder
+    
+    def _apply_precision_policy(self):
+        # (Lógica de precisão permanece a mesma)
+        pass # Placeholder
+
+    def _align(self, dim: int, alignment: int = FRAMES_ALIGNMENT, alignment_rule: str = 'default') -> int:
+        """Aligns a dimension based on a rule."""
+        if alignment_rule == 'n*8+1':
+             return ((dim - 1) // alignment) * alignment + 1
+        return ((dim - 1) // alignment + 1) * alignment
+    
+    def _calculate_aligned_frames(self, duration_s: float, min_frames: int = 1) -> int:
+        num_frames = int(round(duration_s * DEFAULT_FPS))
+        aligned_frames = self._align(num_frames, alignment=FRAMES_ALIGNMENT)
+        return max(aligned_frames, min_frames)
+
+    def _get_random_seed(self) -> int:
+        """Always generates and returns a new random seed."""
+        return random.randint(0, 2**32 - 1)
+
+# ==============================================================================
+# --- INSTANCIAÇÃO SINGLETON ---
+# ==============================================================================
+try:
+    video_generation_service = VideoService()
+    logging.info("Global VideoService orchestrator instance created successfully.")
+except Exception as e:
+    logging.critical(f"Failed to initialize VideoService: {e}", exc_info=True)
+    sys.exit(1)

api/vae_server.py

@@ -0,0 +1,162 @@
+# FILE: api/vae_server.py
+# DESCRIPTION: A dedicated, "hot" VAE service specialist.
+# It loads the VAE model onto a dedicated GPU and keeps it in memory
+# to handle all encoding and decoding requests with minimal latency.
+
+import os
+import sys
+import time
+import logging
+from pathlib import Path
+from typing import List, Union, Tuple
+
+import torch
+import numpy as np
+from PIL import Image
+
+from api.ltx_pool_manager import LatentConditioningItem
+from api.gpu_manager import gpu_manager
+
+
+# --- Importações da Arquitetura e do LTX ---
+try:
+    # Adiciona o path para as bibliotecas do LTX
+    LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
+    if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
+        sys.path.insert(0, str(LTX_VIDEO_REPO_DIR.resolve()))
+    
+    from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
+    from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode
+except ImportError as e:
+    raise ImportError(f"A crucial import failed for VaeServer. Check dependencies. Error: {e}")
+
+
+class VaeServer:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        if self._initialized: return
+        
+        logging.info("⚙️ Initializing VaeServer Singleton...")
+        t0 = time.time()
+        
+        # 1. Obter o dispositivo VAE dedicado do gerenciador central
+        self.device = gpu_manager.get_ltx_vae_device()
+        
+        # 2. Carregar o modelo VAE do checkpoint do LTX
+        # Assumimos que o setup.py já baixou os modelos.
+        try:
+            from api.ltx_pool_manager import ltx_pool_manager
+            # Reutiliza a configuração e o pipeline já carregados pelo LTX Pool Manager
+            # para garantir que estamos usando o mesmo VAE.
+            self.vae = ltx_pool_manager.get_pipeline().vae
+        except Exception as e:
+            logging.critical(f"Failed to get VAE from LTXPoolManager. Is it initialized first? Error: {e}", exc_info=True)
+            raise
+
+        # 3. Garante que o VAE está no dispositivo correto e em modo de avaliação
+        self.vae.to(self.device)
+        self.vae.eval()
+        self.dtype = self.vae.dtype
+        
+        self._initialized = True
+        logging.info(f"✅ VaeServer ready. VAE model is 'hot' on {self.device} with dtype {self.dtype}. Startup time: {time.time() - t0:.2f}s")
+
+    def _cleanup_gpu(self):
+        """Limpa a VRAM da GPU do VAE."""
+        if torch.cuda.is_available():
+            with torch.cuda.device(self.device):
+                torch.cuda.empty_cache()
+
+    def _preprocess_input(self, item: Union[Image.Image, torch.Tensor], target_resolution: Tuple[int, int]) -> torch.Tensor:
+        """Prepara uma imagem PIL ou um tensor para o formato de pixel que o VAE espera."""
+        if isinstance(item, Image.Image):
+            from PIL import ImageOps
+            img = item.convert("RGB")
+            # Redimensiona mantendo a proporção e cortando o excesso
+            processed_img = ImageOps.fit(img, target_resolution, Image.Resampling.LANCZOS)
+            image_np = np.array(processed_img).astype(np.float32) / 255.0
+            tensor = torch.from_numpy(image_np).permute(2, 0, 1) # HWC -> CHW
+        elif isinstance(item, torch.Tensor):
+            # Se já for um tensor, apenas garante que está no formato CHW
+            if item.ndim == 4 and item.shape[0] == 1: # Remove dimensão de batch se houver
+                tensor = item.squeeze(0)
+            elif item.ndim == 3:
+                tensor = item
+            else:
+                raise ValueError(f"Input tensor must have 3 or 4 dimensions (CHW or BCHW), but got {item.ndim}")
+        else:
+            raise TypeError(f"Input must be a PIL Image or a torch.Tensor, but got {type(item)}")
+
+        # Converte para 5D (B, C, F, H, W) e normaliza para [-1, 1]
+        tensor_5d = tensor.unsqueeze(0).unsqueeze(2) # Adiciona B=1 e F=1
+        return (tensor_5d * 2.0) - 1.0
+
+    @torch.no_grad()
+    def generate_conditioning_items(
+        self,
+        media_items: List[Union[Image.Image, torch.Tensor]],
+        target_frames: List[int],
+        strengths: List[float],
+        target_resolution: Tuple[int, int]
+    ) -> List[LatentConditioningItem]:
+        """
+        [FUNÇÃO PRINCIPAL]
+        Converte uma lista de imagens (PIL ou tensores de pixel) em uma lista de
+        LatentConditioningItem, pronta para ser usada pelo pipeline LTX corrigido.
+        """
+        t0 = time.time()
+        logging.info(f"Generating {len(media_items)} latent conditioning items...")
+        
+        if not (len(media_items) == len(target_frames) == len(strengths)):
+            raise ValueError("As listas de media_items, target_frames e strengths devem ter o mesmo tamanho.")
+        
+        conditioning_items = []
+        try:
+            for item, frame, strength in zip(media_items, target_frames, strengths):
+                # 1. Prepara a imagem/tensor para o formato de pixel correto
+                pixel_tensor = self._preprocess_input(item, target_resolution)
+                
+                # 2. Move o tensor de pixel para a GPU do VAE e encoda para latente
+                pixel_tensor_gpu = pixel_tensor.to(self.device, dtype=self.dtype)
+                latents = vae_encode(pixel_tensor_gpu, self.vae, vae_per_channel_normalize=True)
+                
+                # 3. Cria o LatentConditioningItem com o latente (movido para CPU para evitar manter na VRAM)
+                conditioning_items.append(LatentConditioningItem(latents.cpu(), frame, strength))
+
+            logging.info(f"Generated {len(conditioning_items)} items in {time.time() - t0:.2f}s.")
+            return conditioning_items
+        finally:
+            self._cleanup_gpu()
+
+    @torch.no_grad()
+    def decode_to_pixels(self, latent_tensor: torch.Tensor, decode_timestep: float = 0.05) -> torch.Tensor:
+        """Decodifica um tensor latente para um tensor de pixels na CPU."""
+        t0 = time.time()
+        try:
+            latent_tensor_gpu = latent_tensor.to(self.device, dtype=self.dtype)
+            num_items_in_batch = latent_tensor_gpu.shape[0]
+            timestep_tensor = torch.tensor([decode_timestep] * num_items_in_batch, device=self.device, dtype=self.dtype)
+            
+            pixels = vae_decode(
+                latent_tensor_gpu, self.vae, is_video=True, 
+                timestep=timestep_tensor, vae_per_channel_normalize=True
+            )
+            logging.info(f"Decoded latents with shape {latent_tensor.shape} in {time.time() - t0:.2f}s.")
+            return pixels.cpu() # Retorna na CPU
+        finally:
+            self._cleanup_gpu()
+
+# --- Instância Singleton ---
+# A inicialização ocorre quando o módulo é importado pela primeira vez.
+try:
+    vae_server_singleton = VaeServer()
+except Exception as e:
+    logging.critical("CRITICAL: Failed to initialize VaeServer singleton.", exc_info=True)
+    vae_server_singleton = None