Create ltx_server.py

services/ltx_server.py

@@ -0,0 +1,135 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Optional, Tuple
+
+from huggingface_hub import snapshot_download
+
+class LTXServer:
+    """
+    Gerencia o setup e a execução da inferência para o LTX-Video Q8.
+    - Baixa o repositório e os modelos necessários.
+    - Executa o script 'inference.py' como um subprocesso.
+    """
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(LTXServer, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        # Evita reinicialização
+        if hasattr(self, '_initialized') and self._initialized:
+            return
+        
+        print("🚀 LTXServer (Q8) inicializando e preparando o ambiente...")
+        
+        # Define os caminhos
+        self.LTX_REPO_DIR = Path(os.getenv("LTX_REPO_DIR", "/data/LTX-Video"))
+        self.LTX_CKPT_DIR = Path(os.getenv("LTX_CKPT_DIR", "/data/ckpt/ltxvideo_q8"))
+        self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs/ltx"))
+        self.HF_HOME_CACHE = Path(os.getenv("HF_HOME", "/data/.cache/huggingface"))
+        
+        self.REPO_URL = "https://github.com/KONAKONA666/LTX-Video.git"
+        self.MODEL_REPO_ID = "konakona/ltxvideo_q8"
+        
+        # Garante que os diretórios existam
+        for p in [self.LTX_REPO_DIR.parent, self.LTX_CKPT_DIR, self.OUTPUT_ROOT, self.HF_HOME_CACHE]:
+            p.mkdir(parents=True, exist_ok=True)
+            
+        self.setup_dependencies()
+        self._initialized = True
+        print("✅ LTXServer (Q8) pronto.")
+
+    def setup_dependencies(self):
+        """Orquestra o setup: clona o repo e baixa os modelos."""
+        self._ensure_repo()
+        self._ensure_model()
+
+    def _ensure_repo(self) -> None:
+        """Clona o repositório do LTX-Video se ele não existir."""
+        if not (self.LTX_REPO_DIR / ".git").exists():
+            print(f"[LTXServer] Clonando repositório para {self.LTX_REPO_DIR}...")
+            subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.LTX_REPO_DIR)], check=True)
+        else:
+            print("[LTXServer] Repositório LTX-Video já existe.")
+
+    def _ensure_model(self) -> None:
+        """Baixa os checkpoints do modelo Q8, utilizando o cache."""
+        print(f"[LTXServer] Verificando checkpoints em {self.LTX_CKPT_DIR}...")
+        try:
+            snapshot_download(
+                repo_id=self.MODEL_REPO_ID,
+                local_dir=str(self.LTX_CKPT_DIR),
+                cache_dir=str(self.HF_HOME_CACHE),
+                local_dir_use_symlinks=False,
+                repo_type='model',
+                token=os.getenv("HF_TOKEN")
+            )
+            print("[LTXServer] Checkpoints Q8 prontos.")
+        except Exception as e:
+            print(f"[LTXServer] ERRO durante o snapshot_download: {e}")
+            raise
+
+    def run_inference(
+        self,
+        prompt: str,
+        image_path: str,
+        height: int,
+        width: int,
+        num_frames: int,
+        seed: int
+    ) -> str:
+        """Executa a inferência img2vid como um subprocesso."""
+        if not Path(image_path).exists():
+            raise FileNotFoundError(f"Arquivo de imagem de entrada não encontrado: {image_path}")
+
+        script_path = self.LTX_REPO_DIR / "inference.py"
+        if not script_path.exists():
+            raise FileNotFoundError(f"Script de inferência não encontrado em: {script_path}")
+
+        # Cria um diretório de saída único para este job
+        job_output_dir = self.OUTPUT_ROOT / f"run_{int(time.time())}_{os.urandom(4).hex()}"
+        job_output_dir.mkdir(parents=True)
+
+        cmd = [
+            "python", str(script_path),
+            "--ckpt_dir", str(self.LTX_CKPT_DIR),
+            "--output_dir", str(job_output_dir), # Adiciona o diretório de saída
+            "--low_vram",
+            "--transformer_type=q8_kernels",
+            "--prompt", prompt,
+            "--input_image_path", image_path,
+            "--height", str(height),
+            "--width", str(width),
+            "--num_frames", str(num_frames),
+            "--seed", str(seed),
+        ]
+
+        print("[LTXServer] Comando de execução:", " ".join(cmd))
+        
+        try:
+            # O 'cwd' é crucial para que o script encontre seus próprios módulos internos
+            subprocess.run(
+                cmd, 
+                cwd=str(self.LTX_REPO_DIR),
+                check=True, 
+                env=os.environ.copy(),
+                stdout=sys.stdout,
+                stderr=sys.stderr
+            )
+        except Exception as e:
+            print(f"[LTXServer] Erro na execução da inferência: {e}")
+            raise
+            
+        # Encontra o vídeo gerado no diretório de saída
+        output_videos = sorted(job_output_dir.glob("*.mp4"))
+        if not output_videos:
+            raise FileNotFoundError(f"Nenhum vídeo foi gerado no diretório de saída: {job_output_dir}")
+            
+        return str(output_videos[0])
+
+# Instância Singleton que será usada pelo app Gradio
+ltx_server_singleton = LTXServer()