Create ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -0,0 +1,452 @@
+# ltx_server_refactored.py — VideoService (Modular Version)
+
+# --- 0. WARNINGS E AMBIENTE ---
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", message=".*")
+from huggingface_hub import logging
+logging.set_verbosity_error()
+logging.set_verbosity_warning()
+logging.set_verbosity_info()
+logging.set_verbosity_debug()
+LTXV_DEBUG=1
+LTXV_FRAME_LOG_EVERY=8
+import os, subprocess, shlex, tempfile
+import torch
+import json
+import numpy as np
+import random
+import os
+import shlex
+import yaml
+from typing import List, Dict
+from pathlib import Path
+import imageio
+from PIL import Image
+import tempfile
+from huggingface_hub import hf_hub_download
+import sys
+import subprocess
+import gc
+import shutil
+import contextlib
+import time
+import traceback
+from einops import rearrange
+import torch.nn.functional as F
+from managers.vae_manager import vae_manager_singleton
+from tools.video_encode_tool import video_encode_tool_singleton
+DEPS_DIR = Path("/data")
+LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
+
+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)
+        
+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}")
+
+def _query_gpu_processes_via_nvml(device_index: int) -> List[Dict]:
+    try:
+        import psutil
+        import pynvml as nvml
+        nvml.nvmlInit()
+        handle = nvml.nvmlDeviceGetHandleByIndex(device_index)
+        try:
+            procs = nvml.nvmlDeviceGetComputeRunningProcesses_v3(handle)
+        except Exception:
+            procs = nvml.nvmlDeviceGetComputeRunningProcesses(handle)
+        results = []
+        for p in procs:
+            pid = int(p.pid)
+            used_mb = None
+            try:
+                if getattr(p, "usedGpuMemory", None) is not None and p.usedGpuMemory not in (0,):
+                    used_mb = max(0, int(p.usedGpuMemory) // (1024 * 1024))
+            except Exception:
+                used_mb = None
+            name = "unknown"
+            user = "unknown"
+            try:
+                import psutil
+                pr = psutil.Process(pid)
+                name = pr.name()
+                user = pr.username()
+            except Exception:
+                pass
+            results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
+        nvml.nvmlShutdown()
+        return results
+    except Exception:
+        return []
+
+def _query_gpu_processes_via_nvidiasmi(device_index: int) -> List[Dict]:
+    cmd = f"nvidia-smi -i {device_index} --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits"
+    try:
+        out = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT, text=True, timeout=2.0)
+    except Exception:
+        return []
+    results = []
+    for line in out.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) >= 3:
+            try:
+                pid = int(parts[0]); name = parts[1]; used_mb = int(parts[2])
+                user = "unknown"
+                try:
+                    import psutil
+                    pr = psutil.Process(pid)
+                    user = pr.username()
+                except Exception:
+                    pass
+                results.append({"pid": pid, "name": name, "user": user, "used_mb": used_mb})
+            except Exception:
+                continue
+    return results
+
+def calculate_padding(orig_h, orig_w, target_h, target_w):
+    pad_h = target_h - orig_h
+    pad_w = target_w - orig_w
+    pad_top = pad_h // 2
+    pad_bottom = pad_h - pad_top
+    pad_left = pad_w // 2
+    pad_right = pad_w - pad_left
+    return (pad_left, pad_right, pad_top, pad_bottom)
+
+def calculate_new_dimensions(orig_w, orig_h, divisor=8):
+    if orig_w == 0 or orig_h == 0:
+        return 512, 512
+    if orig_w >= orig_h:
+        aspect_ratio = orig_w / orig_h
+        new_h = 512
+        new_w = new_h * aspect_ratio
+    else:
+        aspect_ratio = orig_h / orig_w
+        new_w = 512
+        new_h = new_w * aspect_ratio
+    final_w = int(round(new_w / divisor)) * divisor
+    final_h = int(round(new_h / divisor)) * divisor
+    final_w = max(divisor, final_w)
+    final_h = max(divisor, final_h)
+    print(f"[Dimension Calc] Original: {orig_w}x{orig_h} -> Calculado: {new_w:.0f}x{new_h:.0f} -> Final (divisível por {divisor}): {final_w}x{final_h}")
+    return final_h, final_w
+
+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 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")
+
+add_deps_to_path()
+from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline
+from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
+from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
+from ltx_video.pipelines.pipeline_ltx_video import adain_filter_latent
+from api.ltx.inference import (
+    create_ltx_video_pipeline,
+    create_latent_upsampler,
+    load_image_to_tensor_with_resize_and_crop,
+    seed_everething,
+)
+
+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)}")
+
+        vae_manager_singleton.attach_pipeline(
+            self.pipeline,
+            device=self.device,
+            autocast_dtype=self.runtime_autocast_dtype
+        )
+        print(f"[DEBUG] VAE manager conectado: has_vae={hasattr(self.pipeline, 'vae')} device={self.device}")
+
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+            self._log_gpu_memory("Após carregar modelos")
+
+        print(f"[DEBUG] VideoService pronto. boot_time={time.perf_counter()-t0:.3f}s")
+
+    def _log_gpu_memory(self, stage_name: str):
+        if self.device != "cuda":
+            return
+        device_index = torch.cuda.current_device() if torch.cuda.is_available() else 0
+        current_reserved_b = torch.cuda.memory_reserved(device_index)
+        current_reserved_mb = current_reserved_b / (1024 ** 2)
+        total_memory_b = torch.cuda.get_device_properties(device_index).total_memory
+        total_memory_mb = total_memory_b / (1024 ** 2)
+        peak_reserved_mb = torch.cuda.max_memory_reserved(device_index) / (1024 ** 2)
+        delta_mb = current_reserved_mb - getattr(self, "last_memory_reserved_mb", 0.0)
+        processes = _query_gpu_processes_via_nvml(device_index) or _query_gpu_processes_via_nvidiasmi(device_index)
+        print(f"\n--- [LOG GPU] {stage_name} (cuda:{device_index}) ---")
+        print(f"  - Reservado: {current_reserved_mb:.2f} MB / {total_memory_mb:.2f} MB  (Δ={delta_mb:+.2f} MB)")
+        if peak_reserved_mb > getattr(self, "last_memory_reserved_mb", 0.0):
+            print(f"  - Pico reservado (nesta fase): {peak_reserved_mb:.2f} MB")
+        print(_gpu_process_table(processes, os.getpid()), end="")
+        print("--------------------------------------------------\n")
+        self.last_memory_reserved_mb = current_reserved_mb
+
+    def _register_tmp_dir(self, d: str):
+        if d and os.path.isdir(d):
+            self._tmp_dirs.add(d); print(f"[DEBUG] Registrado tmp dir: {d}")
+
+    def _register_tmp_file(self, f: str):
+        if f and os.path.exists(f):
+            self._tmp_files.add(f); print(f"[DEBUG] Registrado tmp file: {f}")
+
+    def finalize(self, keep_paths=None, extra_paths=None, clear_gpu=True):
+        print("[DEBUG] Finalize: iniciando limpeza...")
+        keep = set(keep_paths or []); extras = set(extra_paths or [])
+        removed_files = 0
+        for f in list(self._tmp_files | extras):
+            try:
+                if f not in keep and os.path.isfile(f):
+                    os.remove(f); removed_files += 1; print(f"[DEBUG] Removido arquivo tmp: {f}")
+            except Exception as e:
+                print(f"[DEBUG] Falha removendo arquivo {f}: {e}")
+            finally:
+                self._tmp_files.discard(f)
+        removed_dirs = 0
+        for d in list(self._tmp_dirs):
+            try:
+                if d not in keep and os.path.isdir(d):
+                    shutil.rmtree(d, ignore_errors=True); removed_dirs += 1; print(f"[DEBUG] Removido diretório tmp: {d}")
+            except Exception as e:
+                print(f"[DEBUG] Falha removendo diretório {d}: {e}")
+            finally:
+                self._tmp_dirs.discard(d)
+        print(f"[DEBUG] Finalize: arquivos removidos={removed_files}, dirs removidos={removed_dirs}")
+        gc.collect()
+        try:
+            if clear_gpu and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                try:
+                    torch.cuda.ipc_collect()
+                except Exception:
+                    pass
+        except Exception as e:
+            print(f"[DEBUG] Finalize: limpeza GPU falhou: {e}")
+        try:
+            self._log_gpu_memory("Após finalize")
+        except Exception as e:
+            print(f"[DEBUG] Log GPU pós-finalize falhou: {e}")
+
+    def _load_config(self):
+        base = LTX_VIDEO_REPO_DIR / "configs"
+        config_path = base / "ltxv-13b-0.9.8-distilled-fp8.yaml"
+        print(f"[DEBUG] Carregando config: {config_path}")
+        with open(config_path, "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"])
+        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"])
+        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
+
+    @torch.no_grad()
+    def _upsample_latents_internal(self, latents: torch.Tensor) -> torch.Tensor:
+        if not self.latent_upsampler:
+            raise ValueError("Latent Upsampler não está carregado.")
+        self.latent_upsampler.to(self.device)
+        self.pipeline.vae.to(self.device)
+        print(f"[DEBUG-UPSAMPLE] Shape de entrada: {tuple(latents.shape)}")
+        latents_unnormalized = un_normalize_latents(latents, self.pipeline.vae, vae_per_channel_normalize=True)
+        upsampled_latents = self.latent_upsampler(latents_unnormalized)
+        upsampled_latents_normalized = normalize_latents(upsampled_latents, self.pipeline.vae, vae_per_channel_normalize=True)
+        print(f"[DEBUG-UPSAMPLE] Shape de saída: {tuple(upsampled_latents_normalized.shape)}")
+        return upsampled_latents_normalized
+        
+    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 in ["float8_e4m3fn", "bfloat16"]:
+            self.runtime_autocast_dtype = torch.bfloat16
+        elif prec == "mixed_precision":
+            self.runtime_autocast_dtype = torch.float16
+
+    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)
+        print(f"[DEBUG] Cond shape={tuple(out.shape)} dtype={out.dtype} device={out.device}")
+        return out
+
+    def _concat_mp4s_no_reencode(self, mp4_list: List[str], out_path: str):
+        if not mp4_list:
+            raise ValueError("A lista de MP4s para concatenar está vazia.")
+        if len(mp4_list) == 1:
+            shutil.move(mp4_list[0], out_path)
+            print(f"[DEBUG] Apenas um vídeo, movido para: {out_path}")
+            return
+            
+        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".txt") as f:
+            for mp4 in mp4_list:
+                f.write(f"file '{os.path.abspath(mp4)}'\n")
+            list_path = f.name
+    
+        cmd = f"ffmpeg -y -f concat -safe 0 -i {list_path} -c copy {out_path}"
+        print(f"[DEBUG] Concat: {cmd}")
+    
+        try:
+            subprocess.check_call(shlex.split(cmd))
+        finally:
+            os.remove(list_path)
+
+    def _save_and_log_video(self, pixel_tensor, base_filename, fps, temp_dir, results_dir, used_seed, progress_callback=None):
+        """Função auxiliar para salvar um tensor de pixels em um arquivo MP4."""
+        output_path = os.path.join(temp_dir, f"{base_filename}_{used_seed}.mp4")
+        
+        video_encode_tool_singleton.save_video_from_tensor(
+            pixel_tensor, output_path, fps=fps, progress_callback=progress_callback
+        )
+        
+        final_path = os.path.join(results_dir, f"{base_filename}_{used_seed}.mp4")
+        shutil.move(output_path, final_path)
+        print(f"[DEBUG] Vídeo salvo em: {final_path}")
+        return final_path
+
+    # ==============================================================================
+    # --- NOVAS FUNÇÕES MODULARES ---
+    # ==============================================================================
+
+    def prepare_condition_items(self, items_list: List, height: int, width: int, num_frames: int):
+        """
+        Prepara a lista de tensores de condicionamento a partir de uma lista de imagens ou tensores.
+        Formato da lista de entrada: [[media_path_ou_tensor, frame_alvo, peso], ...]
+        """
+        if not items_list:
+            return []
+
+        height_padded = ((height - 1) // 8 + 1) * 8
+        width_padded = ((width - 1) // 8 + 1) * 8
+        padding_values = calculate_padding(height, width, height_padded, width_padded)
+        
+        conditioning_items = []
+        print("\n--- Preparando Itens de Condicionamento ---")
+        for item in items_list:
+            media, frame, weight = item
+            
+            if isinstance(media, str): 
+                print(f"  - Carregando imagem: {media} para o frame {frame}")
+                tensor = self._prepare_conditioning_tensor(media, height, width, padding_values)
+            elif isinstance(media, torch.Tensor):
+                print(f"  - Usando tensor fornecido para o frame {frame}")
+                tensor = media.to(self.device, dtype=self.runtime_autocast_dtype)
+            else:
+                warnings.warn(f"Tipo de item desconhecido: {type(media)}. Ignorando.")
+                continue
+
+            safe_frame = max(0, min(int(frame), num_frames - 1))
+            conditioning_items.append(ConditioningItem(tensor, safe_frame, float(weight)))
+
+        print(f"Total de itens de condicionamento preparados: {len(conditioning_items)}")
+        return conditioning_items
+
+    def generate_low(self, prompt, negative_prompt, height, width, duration, guidance_scale, seed, conditioning_items=None):
+        """
+        Gera um vídeo em baixa resolução (primeiro passe).
+        Retorna: (caminho_do_video_mp4, caminho_do_tensor_cpu, seed_usado)
+        """
+        print("\n--- INICIANDO ETAPA 1: GERAÇÃO EM BAIXA RESOLUÇÃO ---")
+        self._log_gpu_memory("Início da Geração Low-Res")
+        
+        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
+        seed_everething(used_seed)
+        
+        FPS = 24.0
+        target_frames = round(duration * FPS)
+        actual_num_frames = max(9, int(round((target_frames - 1) / 8.0) * 8 + 1))
+        
+        height_padded = ((height - 1) // 8 + 1) * 8
+        width_padded = ((width - 1) // 8 + 1) * 8
+        generator = torch.Generator(device=self.device).manual_seed(used_seed)
+        
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_low_"); self._register_tmp_dir(temp_dir)
+        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
+        
+        downscale_factor = self