# FILE: api/ltx/ltx_aduc_pipeline.py # DESCRIPTION: Final high-level orchestrator for LTX-Video generation. # This version acts as a client to the specialized managers (LTX, VAE), # focusing solely on the business logic of video generation workflows. 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 random import torch import yaml import numpy as np from PIL import Image from api.ltx.ltx_utils import seed_everything from utils.debug_utils import log_function_io from managers.gpu_manager import gpu_manager from api.ltx.ltx_aduc_manager import ltx_aduc_manager, LatentConditioningItem from api.ltx.vae_aduc_pipeline import vae_aduc_pipeline from tools.video_encode_tool import video_encode_tool_singleton # ============================================================================== # --- 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 = logging.DEBUG 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 # Garante que a biblioteca LTX-Video seja importável repo_path = str(LTX_VIDEO_REPO_DIR.resolve()) if repo_path not in sys.path: sys.path.insert(0, repo_path) # ============================================================================== # --- CLASSE DE SERVIÇO (O ORQUESTRADOR) --- # ============================================================================== class LtxAducPipeline: """ Orchestrates the high-level logic of video generation, delegating all low-level tasks to specialized managers and utility modules. """ @log_function_io def __init__(self): t0 = time.time() logging.info("Initializing VideoService Orchestrator...") if ltx_aduc_manager is None or vae_aduc_pipeline is None: raise RuntimeError("A required manager (LTX or VAE) failed to initialize. Aborting.") self.pipeline = ltx_aduc_manager.get_pipeline() self.main_device = self.pipeline.device self.vae_device = self.pipeline.vae.device self.config = ltx_aduc_manager.config self._apply_precision_policy() logging.info(f"VideoService ready. Using Main: {self.main_device}, VAE: {self.vae_device}. Startup time: {time.time() - t0:.2f}s") def finalize(self): """Cleans up GPU memory after a generation task.""" gc.collect() if torch.cuda.is_available(): with torch.cuda.device(self.main_device): torch.cuda.empty_cache() with torch.cuda.device(self.vae_device): 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 video from a list of prompts and raw 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 list cannot be empty.") 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("Delegating to VaeServer to prepare initial conditioning items...") initial_conditions = vae_aduc_pipeline.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]}...'") current_frames_base = frames_per_chunk if i < num_chunks - 1 else total_frames - ((num_chunks - 1) * frames_per_chunk) 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[LatentConditioningItem]): """ Logs detailed information about a list of ConditioningItem objects. This is a dedicated debug helper function. """ # Só imprime o log se o nível de logging for DEBUG if logging.getLogger().isEnabledFor(logging.INFO): log_str = ["\n" + "="*10 + " INFO: Conditioning Items " + "="*10] if not items: log_str.append(" -> Lista de conditioning_items está vazia.") else: for i, item in enumerate(items): if hasattr(item, 'media_item') and isinstance(item.media_item, torch.Tensor): t = item.media_item log_str.append( f" -> Item [{i}]: " f"Tensor(shape={list(t.shape)}, " f"device='{t.device}', " f"dtype={t.dtype}), " f"Target Frame = {item.media_frame_number}, " f"Strength = {item.conditioning_strength:.2f}" ) else: tt = str(itemvalue) log_str.append(f" -> Item [{i}]: Não contém um tensor válido.") log_str.append(f" {tt[:70]}") log_str.append("="*40 + "\n") # Usa o logger de debug para imprimir a mensagem completa logging.info("\n".join(log_str)) @log_function_io def _generate_single_chunk_low(self, **kwargs) -> Optional[torch.Tensor]: """[WORKER] Calls the pipeline to generate a single chunk of latents.""" height_padded, width_padded = (self._align(d) for d in (kwargs['height'], kwargs['width'])) downscale_factor = self.config.get("downscale_factor", 0.6666666) vae_scale_factor = self.pipeline.vae_scale_factor downscaled_height = self._align(int(height_padded * downscale_factor), vae_scale_factor) downscaled_width = self._align(int(width_padded * downscale_factor), vae_scale_factor) # 1. Começa com a configuração padrão first_pass_config = self.config.get("first_pass", {}).copy() # 2. Aplica os overrides da UI, se existirem if kwargs.get("ltx_configs_override"): self._apply_ui_overrides(first_pass_config, kwargs.get("ltx_configs_override")) # 3. Monta o dicionário de argumentos SEM conditioning_items primeiro pipeline_kwargs = { "prompt": kwargs['prompt'], "negative_prompt": kwargs['negative_prompt'], "height": downscaled_height, "width": downscaled_width, "num_frames": kwargs['num_frames'], "frame_rate": int(DEFAULT_FPS), "generator": torch.Generator(device=self.main_device).manual_seed(kwargs['seed']), "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, **first_pass_config } # --- Bloco de Logging para Depuração --- # 4. Loga os argumentos do pipeline (sem os tensores de condição) logging.info(f"\n[Info] Pipeline Arguments (BASE):\n {json.dumps(pipeline_kwargs, indent=2, default=str)}\n") # Loga os conditioning_items separadamente com a nossa função helper conditioning_items_list = kwargs.get('conditioning_items') self._log_conditioning_items(conditioning_items_list) pipeline_kwargs['conditioning_items'] = conditioning_items_list with torch.autocast(device_type=self.main_device.type, dtype=self.runtime_autocast_dtype, enabled="cuda" in self.main_device.type): latents_raw = self.pipeline(**pipeline_kwargs).images return latents_raw.to(self.main_device) @log_function_io def _finalize_generation(self, final_latents: torch.Tensor, base_filename: str, seed: int) -> Tuple[str, str]: """Delegates final decoding and encoding to specialist services.""" logging.info("Finalizing generation: decoding latents and encoding 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_aduc_pipeline.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.""" # Override step counts for key in ["num_inference_steps", "skip_initial_inference_steps", "skip_final_inference_steps"]: ui_value = overrides.get(key) if ui_value and ui_value > 0: config_dict[key] = ui_value logging.info(f"Override: '{key}' set to {ui_value} by UI.") def _save_and_log_video(self, pixel_tensor: torch.Tensor, base_filename: str) -> Path: with tempfile.TemporaryDirectory() as temp_dir: temp_path = os.path.join(temp_dir, f"{base_filename}.mp4") video_encode_tool_singleton.save_video_from_tensor(pixel_tensor, temp_path, fps=DEFAULT_FPS) final_path = RESULTS_DIR / f"{base_filename}.mp4" shutil.move(temp_path, final_path) logging.info(f"Video saved successfully to: {final_path}") return final_path def _apply_precision_policy(self): precision = str(self.config.get("precision", "bfloat16")).lower() if precision in ["float8_e4m3fn", "bfloat16"]: self.runtime_autocast_dtype = torch.bfloat16 elif precision == "mixed_precision": self.runtime_autocast_dtype = torch.float16 else: self.runtime_autocast_dtype = torch.float32 logging.info(f"Runtime precision policy set for autocast: {self.runtime_autocast_dtype}") def _align(self, dim: int, alignment: int = FRAMES_ALIGNMENT, alignment_rule: str = 'default') -> int: 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: return random.randint(0, 2**32 - 1) # ============================================================================== # --- INSTANCIAÇÃO SINGLETON --- # ============================================================================== ltx_aduc_pipeline = LtxAducPipeline() logging.info("Global VideoService orchestrator instance created successfully.")