# ltx_server.py — VideoService (beta 1.0) # Sempre output_type="latent"; no final: VAE (bloco inteiro) → pixels → MP4. # --- 1. IMPORTAÇÕES --- import torch 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 # Singletons do projeto para VAE e Encoder from tools.video_encode_tool import video_encode_tool_singleton from managers.vae_manager import vae_manager_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: 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 _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) 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 inference import ( create_ltx_video_pipeline, create_latent_upsampler, load_image_to_tensor_with_resize_and_crop, seed_everething, calculate_padding, load_media_file, ) from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy # --- 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)}") 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-dev-fp8.yaml.txt", 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.yaml" print(f"[DEBUG] Config fallback: {cfg}") with open(cfg, "r") as file: return yaml.safe_load(file) def _load_models(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}") 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 # --- 6. GERAÇÃO --- 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, # Sempre latent → VAE → MP4 (simples) 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") 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") if mode == "video-to-video" and not input_video_filepath: raise ValueError("O vídeo de entrada é obrigatório para o modo video-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 = 257 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))) print(f"[DEBUG] Frames alvo: {actual_num_frames} (dur={duration}s @ {FPS}fps)") height_padded = ((height - 1) // 32 + 1) * 32 width_padded = ((width - 1) // 32 + 1) * 32 padding_values = calculate_padding(height, width, height_padded, width_padded) print(f"[DEBUG] Dimensões: ({height},{width}) -> pad ({height_padded},{width_padded}); padding={padding_values}") 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)}") # Sempre pedimos latentes (simples) 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.15, "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, } print(f"[DEBUG] output_type={call_kwargs['output_type']} skip_layer_strategy={call_kwargs['skip_layer_strategy']}") if mode == "video-to-video": media = load_media_file( media_path=input_video_filepath, height=height, width=width, max_frames=int(frames_to_use), padding=padding_values, ).to(self.device) call_kwargs["media_items"] = media print(f"[DEBUG] media_items shape={tuple(media.shape)}") latents = None multi_scale_pipeline = None try: if improve_texture: if not self.latent_upsampler: raise ValueError("Upscaler espacial não carregado.") print("[DEBUG] Multi-escala: construindo pipeline...") multi_scale_pipeline = LTXMultiScalePipeline(self.pipeline, self.latent_upsampler) first_pass_args = self.config.get("first_pass", {}).copy() first_pass_args["guidance_scale"] = float(guidance_scale) second_pass_args = self.config.get("second_pass", {}).copy() second_pass_args["guidance_scale"] = float(guidance_scale) multi_scale_call_kwargs = call_kwargs.copy() multi_scale_call_kwargs.update( { "downscale_factor": self.config["downscale_factor"], "first_pass": first_pass_args, "second_pass": second_pass_args, } ) print("[DEBUG] Chamando multi_scale_pipeline...") t_ms = time.perf_counter() ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext() with ctx: result = multi_scale_pipeline(**multi_scale_call_kwargs) print(f"[DEBUG] multi_scale_pipeline tempo={time.perf_counter()-t_ms:.3f}s") # Captura latentes if hasattr(result, "latents"): latents = result.latents elif hasattr(result, "images") and isinstance(result.images, torch.Tensor): latents = result.images else: latents = result print(f"[DEBUG] Latentes (multi-escala): shape={tuple(latents.shape)}") else: single_pass_kwargs = call_kwargs.copy() first_pass_config = self.config.get("first_pass", {}) single_pass_kwargs.update( { "guidance_scale": float(guidance_scale), "stg_scale": first_pass_config.get("stg_scale"), "rescaling_scale": first_pass_config.get("rescaling_scale"), "skip_block_list": first_pass_config.get("skip_block_list"), } ) schedule = first_pass_config.get("timesteps") or first_pass_config.get("guidance_timesteps") if mode == "video-to-video": schedule = [0.7]; print("[INFO] Modo video-to-video (etapa única): timesteps=[0.7]") if isinstance(schedule, (list, tuple)) and len(schedule) > 0: single_pass_kwargs["timesteps"] = schedule single_pass_kwargs["guidance_timesteps"] = schedule print(f"[DEBUG] Single-pass: timesteps_len={len(schedule) if schedule else 0}") print("\n[INFO] Executando pipeline de etapa única...") t_sp = time.perf_counter() ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext() with ctx: result = self.pipeline(**single_pass_kwargs) print(f"[DEBUG] single-pass tempo={time.perf_counter()-t_sp:.3f}s") if hasattr(result, "latents"): latents = result.latents elif hasattr(result, "images") and isinstance(result.images, torch.Tensor): latents = result.images else: latents = result print(f"[DEBUG] Latentes (single-pass): shape={tuple(latents.shape)}") # Staging e escrita MP4 (simples: VAE → pixels → MP4) temp_dir = tempfile.mkdtemp(prefix="ltxv_"); self._register_tmp_dir(temp_dir) results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True) output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4") final_output_path = None pixel_tensor = vae_manager_singleton.decode( latents.to(self.device, non_blocking=True), decode_timestep=float(self.config.get("decode_timestep", 0.05)) ) print("[DEBUG] Decodificando bloco de latentes com VAE → tensor de pixels...") # Se desejar “desocupar” a GPU antes do decode, pode-se mover p/ CPU e limpar: # latents_cpu = latents.detach().to("cpu", non_blocking=True); torch.cuda.empty_cache(); torch.cuda.ipc_collect(); latents = latents_cpu.to(self.device) pixel_tensor = vae_manager_singleton.decode(latents.to(self.device, non_blocking=True)) 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"] ) candidate_final = os.path.join(results_dir, f"output_{used_seed}.mp4") try: shutil.move(output_video_path, candidate_final) final_output_path = candidate_final print(f"[DEBUG] MP4 movido para {final_output_path}") except Exception as e: final_output_path = output_video_path print(f"[DEBUG] Falha no move; usando tmp como final: {e}") self._register_tmp_file(output_video_path) self._log_gpu_memory("Fim da Geração") print(f"[DEBUG] generate() fim ok. total_time={time.perf_counter()-t_all:.3f}s") return final_output_path, 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()