Update services/ltx_server.py

services/ltx_server.py

@@ -3,15 +3,22 @@ import subprocess
 import sys
 import time
 import shutil
+import json  # Importa o módulo json para manipular o arquivo de configuração
 from pathlib import Path
 from typing import Optional, Tuple
 
-# Para o download seletivo, hf_hub_download é mais direto que snapshot_download
 from huggingface_hub import snapshot_download
 
+# Define a raiz da aplicação a partir de variáveis de ambiente para robustez.
 APP_HOME = Path(os.environ.get("APP_HOME", "/app"))
 
 class LTXServer:
+    """
+    Gerencia o setup e a execução da inferência para o LTX-Video Q8.
+    - Baixa os componentes de modelo de dois repositórios diferentes.
+    - Aplica um patch no config.json do VAE para garantir compatibilidade.
+    - Executa o script 'inference.py' como um subprocesso.
+    """
     _instance = None
 
     def __new__(cls, *args, **kwargs):
@@ -20,10 +27,12 @@ class LTXServer:
         return cls._instance
 
     def __init__(self):
-        if hasattr(self, '_initialized') and self._initialized: return
+        if hasattr(self, '_initialized') and self._initialized:
+            return
         
-        print("🚀 LTXServer (Q8) inicializando com download seletivo...")
+        print("🚀 LTXServer (Q8) inicializando e preparando o ambiente...")
         
+        # Define os caminhos principais
         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 = APP_HOME / "outputs" / "ltx"
@@ -33,6 +42,7 @@ class LTXServer:
         self.MODEL_REPO_ID_ORIGINAL = "Lightricks/LTX-Video"
         self.MODEL_REPO_ID_Q8 = "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)
             
@@ -41,20 +51,22 @@ class LTXServer:
         print("✅ LTXServer (Q8) pronto.")
 
     def setup_dependencies(self):
+        """Orquestra o setup: clona o repo, baixa os modelos e aplica o patch."""
         self._ensure_repo()
-        self._ensure_model()
+        self._ensure_model_and_patch_config()
 
     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:
+    def _ensure_model_and_patch_config(self) -> None:
         """
-        Garante que todos os componentes do modelo existam, baixando apenas os
-        subdiretórios necessários de cada repositório.
+        Garante que todos os componentes existam e aplica um patch no config.json
+        do VAE para corresponder ao que o código do fork espera.
         """
         print(f"[LTXServer] Verificando e baixando componentes em {self.LTX_CKPT_DIR}...")
         
@@ -62,12 +74,7 @@ class LTXServer:
         cache_dir = str(self.HF_HOME_CACHE)
         local_dir = str(self.LTX_CKPT_DIR)
 
-        # ====================================================================
-        # <<< NOVA LÓGICA DE DOWNLOAD SELETIVO >>>
-
-        # 1. Baixa os componentes base (VAE, Text Encoder, Scheduler) do repo original.
-        # O padrão "glob" `componente/*` garante que baixemos todo o conteúdo de cada pasta.
-        componentes_base = ["vae/*", "text_encoder/*", "scheduler/*"]
+        # 1. Baixa os componentes base (VAE, Text Encoder, Scheduler) do repo original
         print(f"[LTXServer] Baixando componentes base de '{self.MODEL_REPO_ID_ORIGINAL}'...")
         snapshot_download(
             repo_id=self.MODEL_REPO_ID_ORIGINAL,
@@ -75,12 +82,12 @@ class LTXServer:
             cache_dir=cache_dir,
             repo_type='model',
             token=token,
-            allow_patterns=componentes_base,
+            ignore_patterns=["unet/*", "*.safetensors"],
             resume_download=True,
         )
         print("[LTXServer] Componentes base (VAE, T5, Scheduler) prontos.")
 
-        # 2. Baixa apenas o UNet quantizado (Q8) do repo do konakona.
+        # 2. Baixa apenas o UNet quantizado (Q8) do repo do konakona
         print(f"[LTXServer] Baixando UNet quantizado (Q8) de '{self.MODEL_REPO_ID_Q8}'...")
         snapshot_download(
             repo_id=self.MODEL_REPO_ID_Q8,
@@ -88,18 +95,33 @@ class LTXServer:
             cache_dir=cache_dir,
             repo_type='model',
             token=token,
-            allow_patterns=["unet/*"], # Baixa apenas a pasta unet
+            allow_patterns=["unet/*"],
             resume_download=True,
         )
         print("[LTXServer] UNet quantizado (Q8) pronto.")
-        # ====================================================================
+        
+        # 3. Aplica o patch no config.json do VAE para resolver o AssertionError
+        vae_config_path = self.LTX_CKPT_DIR / "vae" / "config.json"
+        if vae_config_path.exists():
+            print(f"[LTXServer] Aplicando patch de compatibilidade ao '{vae_config_path.name}' do VAE...")
+            with open(vae_config_path, 'r+') as f:
+                config_data = json.load(f)
+                # Adiciona ou sobrescreve a chave para corresponder ao que o código do fork espera
+                config_data["_class_name"] = "CausalVideoAutoencoder"
+                f.seek(0)
+                json.dump(config_data, f, indent=4)
+                f.truncate()
+            print("[LTXServer] Patch do config.json aplicado com sucesso.")
+        else:
+            # Isso seria um erro crítico, pois significa que o download falhou
+            raise FileNotFoundError(f"Não foi possível encontrar {vae_config_path} para aplicar o patch.")
 
         print("[LTXServer] Todos os componentes do modelo foram baixados e mesclados com sucesso.")
 
     def run_inference(self, prompt: str, image_path: str, height: int, width: int, num_frames: int, seed: int) -> str:
-        # ... (O resto da classe permanece o mesmo da resposta anterior)
         script_path = self.LTX_REPO_DIR / "inference.py"
-        if not script_path.exists(): raise FileNotFoundError(f"Script de inferência não encontrado: {script_path}")
+        if not script_path.exists():
+            raise FileNotFoundError(f"Script de inferência não encontrado: {script_path}")
         
         job_output_dir = self.OUTPUT_ROOT / f"run_{int(time.time())}_{os.urandom(4).hex()}"
         job_output_dir.mkdir(parents=True)
@@ -123,8 +145,12 @@ class LTXServer:
         
         try:
             subprocess.run(
-                cmd, cwd=str(self.LTX_REPO_DIR), check=True, 
-                env=os.environ.copy(), stdout=sys.stdout, stderr=sys.stderr
+                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}")