Update api/ltx_server_refactored_complete.py

api/ltx_server_refactored_complete.py

@@ -126,104 +126,171 @@ class VideoService:
     """
 
     def __init__(self):
-        """Initializes the service, loads models, and configures the environment."""
         t0 = time.perf_counter()
-        logging.info("Initializing VideoService...")
-        RESULTS_DIR.mkdir(parents=True, exist_ok=True)
-
-        self.config = self._load_config(DEFAULT_CONFIG_FILE)
-        self._tmp_dirs = set()
+        print("[DEBUG] Inicializando VideoService...")
         
-        self.pipeline, self.latent_upsampler = self._load_models_on_cpu()
-
+        # 1. Obter o dispositivo alvo a partir do gerenciador
+        #    Não definimos `self.device` ainda, apenas guardamos o alvo.
         target_device = gpu_manager.get_ltx_device()
-        self.device = torch.device("cpu") # Default device
-        self.move_to_device(target_device)
+        print(f"[DEBUG] LTX foi alocado para o dispositivo: {target_device}")
 
+        # 2. Carregar a configuração e os modelos (na CPU, como a função _load_models faz)
+        self.config = self._load_config()
+        self.pipeline, self.latent_upsampler = self._load_models()
+
+        # 3. Mover os modelos para o dispositivo alvo e definir `self.device`
+        self.move_to_device(target_device) # Usando a função que já criamos!
+
+        # 4. Configurar o resto dos componentes com o dispositivo correto
         self._apply_precision_policy()
         vae_manager_singleton.attach_pipeline(
             self.pipeline,
-            device=self.device,
+            device=self.device, # Agora `self.device` está correto
             autocast_dtype=self.runtime_autocast_dtype
         )
+        self._tmp_dirs = set()
+        print(f"[DEBUG] VideoService pronto. boot_time={time.perf_counter()-t0:.3f}s")
+    
+    # A função move_to_device que criamos antes é essencial aqui
+    def move_to_device(self, device):
+        """Move os modelos do pipeline para o dispositivo especificado."""
+        print(f"[LTX] Movendo modelos para {device}...")
+        self.device = torch.device(device) # Garante que é um objeto torch.device
+        self.pipeline.to(self.device)
+        if self.latent_upsampler:
+            self.latent_upsampler.to(self.device)
+        print(f"[LTX] Modelos agora estão em {self.device}.")
         
-        logging.info(f"VideoService ready. Startup time: {time.perf_counter()-t0:.2f}s")
+    def move_to_cpu(self):
+        """Move os modelos para a CPU para liberar VRAM."""
+        self.move_to_device(torch.device("cpu"))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
+            
     # ==========================================================================
     # --- LIFECYCLE & MODEL MANAGEMENT ---
     # ==========================================================================
 
-    def _load_config(self, config_path: Path) -> Dict:
-        """Loads the YAML configuration file."""
-        logging.info(f"Loading config from: {config_path}")
+    def _load_config(self):
+        base = LTX_VIDEO_REPO_DIR / "configs"
+        config_path = base / "ltxv-13b-0.9.8-distilled-fp8.yaml"
         with open(config_path, "r") as file:
             return yaml.safe_load(file)
 
-    def _load_models_on_cpu(self) -> Tuple[LTXMultiScalePipeline, Optional[torch.nn.Module]]:
-        """Downloads and loads the pipeline and upsampler checkpoints onto the CPU."""
+    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 [])
+        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}")
+            
+    def _load_models(self):
         t0 = time.perf_counter()
-        
-        logging.info("Downloading main checkpoint...")
+        LTX_REPO = "Lightricks/LTX-Video"
+        print("[DEBUG] Baixando checkpoint principal...")
         distilled_model_path = hf_hub_download(
-            repo_id=LTX_REPO_ID,
+            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"],
-            device="cpu", # Load on CPU first
-            # Pass other config values directly
-            **{k: v for k, v in self.config.items() if k in create_ltx_video_pipeline.__code__.co_varnames}
+            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"):
-            logging.info("Downloading spatial upscaler checkpoint...")
-            spatial_upscaler_path = hf_hub_download(
-                repo_id=LTX_REPO_ID,
-                filename=self.config["spatial_upscaler_model_path"],
-                token=os.getenv("HF_TOKEN")
-            )
-            self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
+            print("[DEBUG] Construindo latent_upsampler...")
             latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
-
-        logging.info(f"Models loaded on CPU in {time.perf_counter()-t0:.2f}s")
+            print("[DEBUG] Upsampler pronto.")
+        print(f"[DEBUG] _load_models() tempo total={time.perf_counter()-t0:.3f}s")
         return pipeline, latent_upsampler
 
-    def move_to_device(self, device_str: str):
-        """Moves all relevant models to the specified device (e.g., 'cuda:0' or 'cpu')."""
-        target_device = torch.device(device_str)
-        if self.device == target_device:
-            logging.info(f"Models are already on the target device: {device_str}")
-            return
-        
-        logging.info(f"Moving models to {device_str}...")
-        self.device = target_device
-        self.pipeline.to(self.device)
-        if self.latent_upsampler:
-            self.latent_upsampler.to(self.device)
-        
-        if device_str == "cpu" and torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            
-        logging.info(f"Models successfully moved to {self.device}.")
+    def _apply_precision_policy(self):
+        prec = str(self.config.get("precision", "")).lower()
+        self.runtime_autocast_dtype = torch.float32
+        if prec in ["float8_e4m3fn", "bfloat16"]:
+            self.runtime_autocast_dtype = torch.bfloat16
+        elif prec == "mixed_precision":
+            self.runtime_autocast_dtype = torch.float16
 
-    def finalize(self, keep_paths: Optional[List[str]] = None):
-        """Cleans up GPU memory and temporary directories."""
-        logging.debug("Finalizing resources...")
-        gc.collect()
-        if torch.cuda.is_available():
+    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}")
+
+    @torch.no_grad()
+    def _upsample_latents_internal(self, latents: torch.Tensor) -> torch.Tensor:
+        try:
+            if not self.latent_upsampler:
+                raise ValueError("Latent Upsampler não está carregado.")
+            latents_unnormalized = un_normalize_latents(latents, self.pipeline.vae, vae_per_channel_normalize=True)
+            upsampled_latents = self.latent_upsampler(latents_unnormalized)
+            return normalize_latents(upsampled_latents, self.pipeline.vae, vae_per_channel_normalize=True)
+        except Exception as e:
+            pass
+        finally:
             torch.cuda.empty_cache()
-            try:
-                torch.cuda.ipc_collect()
-            except Exception:
-                pass
-        
-        # Optional: Clean up temporary directories if needed (logic can be added here)
+            torch.cuda.ipc_collect()
+            self.finalize(keep_paths=[])
+
+    def _prepare_conditioning_tensor(self, filepath, height, width, padding_values):
+        tensor = load_image_to_tensor_with_resize_and_crop(filepath, height, width)
+        tensor = torch.nn.functional.pad(tensor, padding_values)
+        log_tensor_info(tensor, f"_prepare_conditioning_tensor")
+        return tensor.to(self.device, dtype=self.runtime_autocast_dtype)
 
+            
+    def _save_and_log_video(self, pixel_tensor, base_filename, fps, temp_dir, results_dir, used_seed, progress_callback=None):
+        output_path = os.path.join(temp_dir, f"{base_filename}_.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}_.mp4")
+        shutil.move(output_path, final_path)
+        print(f"[DEBUG] Vídeo salvo em: {final_path}")
+        return final_path
+
+    def _load_tensor(self, caminho):
+        # Se já é um tensor, retorna diretamente
+        if isinstance(caminho, torch.Tensor):
+            return caminho
+        # Se é bytes, carrega do buffer
+        if isinstance(caminho, (bytes, bytearray)):
+            return torch.load(io.BytesIO(caminho))
+        # Caso contrário, assume que é um caminho de arquivo
+        return torch.load(caminho
 
     # ==========================================================================
     # --- PUBLIC ORCHESTRATORS ---