Update api/ltx_server_refactored_complete.py

api/ltx_server_refactored_complete.py

@@ -1,7 +1,7 @@
 # FILE: api/ltx_server_refactored_complete.py
-# DESCRIPTION: Final high-level orchestrator for LTX-Video generation.
-# This version features a unified generation workflow, random seed generation,
-# delegation to specialized modules, and advanced debugging capabilities.
+# DESCRIPTION: Final orchestrator for LTX-Video generation.
+# This version includes the fix for the narrative generation overlap bug and
+# consolidates all previous refactoring and debugging improvements.
 
 import gc
 import json
@@ -36,7 +36,7 @@ LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
 RESULTS_DIR = Path("/app/output")
 DEFAULT_FPS = 24.0
 FRAMES_ALIGNMENT = 8
-LTX_REPO_ID = "Lightricks/LTX-Video" # Repositório de onde os modelos são baixados
+LTX_REPO_ID = "Lightricks/LTX-Video"
 
 # Garante que a biblioteca LTX-Video seja importável
 def add_deps_to_path():
@@ -45,7 +45,7 @@ def add_deps_to_path():
         sys.path.insert(0, repo_path)
         logging.info(f"[ltx_server] LTX-Video repository added to sys.path: {repo_path}")
 
-#add_deps_to_path()
+add_deps_to_path()
 
 # --- Módulos da nossa Arquitetura ---
 try:
@@ -99,7 +99,7 @@ class VideoService:
         logging.info(f"LTX allocated to devices: Main='{target_main_device_str}', VAE='{target_vae_device_str}'")
 
         self.config = self._load_config()
-        self._resolve_model_paths_from_cache() # Etapa crítica para encontrar os modelos
+        self._resolve_model_paths_from_cache()
 
         self.pipeline, self.latent_upsampler = build_ltx_pipeline_on_cpu(self.config)
 
@@ -119,31 +119,16 @@ class VideoService:
             return yaml.safe_load(file)
 
     def _resolve_model_paths_from_cache(self):
-        """
-        Uses hf_hub_download to find the absolute paths to model files in the cache,
-        updating the in-memory config. This makes the app resilient to cache structure.
-        """
+        """Finds the absolute paths to model files in the cache and updates the in-memory config."""
         logging.info("Resolving model paths from Hugging Face cache...")
         cache_dir = os.environ.get("HF_HOME")
         try:
-            # Resolve o caminho do checkpoint principal
-            main_ckpt_filename = self.config["checkpoint_path"]
-            main_ckpt_path = hf_hub_download(
-                repo_id=LTX_REPO_ID,
-                filename=main_ckpt_filename,
-                cache_dir=cache_dir
-            )
+            main_ckpt_path = hf_hub_download(repo_id=LTX_REPO_ID, filename=self.config["checkpoint_path"], cache_dir=cache_dir)
             self.config["checkpoint_path"] = main_ckpt_path
             logging.info(f"  -> Main checkpoint resolved to: {main_ckpt_path}")
 
-            # Resolve o caminho do upsampler, se existir
             if self.config.get("spatial_upscaler_model_path"):
-                upscaler_filename = self.config["spatial_upscaler_model_path"]
-                upscaler_path = hf_hub_download(
-                    repo_id=LTX_REPO_ID,
-                    filename=upscaler_filename,
-                    cache_dir=cache_dir
-                )
+                upscaler_path = hf_hub_download(repo_id=LTX_REPO_ID, filename=self.config["spatial_upscaler_model_path"], cache_dir=cache_dir)
                 self.config["spatial_upscaler_model_path"] = upscaler_path
                 logging.info(f"  -> Spatial upscaler resolved to: {upscaler_path}")
         except Exception as e:
@@ -201,7 +186,11 @@ class VideoService:
         num_chunks = len(prompt_list)
         total_frames = self._calculate_aligned_frames(kwargs.get("duration", 4.0))
         frames_per_chunk = max(FRAMES_ALIGNMENT, (total_frames // num_chunks // FRAMES_ALIGNMENT) * FRAMES_ALIGNMENT)
-        overlap_frames = self.config.get("overlap_frames", 8) if is_narrative else 0
+        
+        # Overlap must be N*8+1 frames. 9 is the smallest practical value.
+        overlap_frames = 9 if is_narrative else 0
+        if is_narrative:
+            logging.info(f"Narrative mode: Using overlap of {overlap_frames} frames between chunks.")
         
         temp_latent_paths = []
         overlap_condition_item = None
@@ -210,14 +199,16 @@ class VideoService:
             for i, chunk_prompt in enumerate(prompt_list):
                 logging.info(f"Processing scene {i+1}/{num_chunks}: '{chunk_prompt[:50]}...'")
                 
-                if i == num_chunks - 1:
-                    processed_frames = (num_chunks - 1) * frames_per_chunk
-                    current_frames = total_frames - processed_frames
-                else:
-                    current_frames = frames_per_chunk
-                
-                if i > 0: current_frames += overlap_frames
+                if i < num_chunks - 1:
+                    current_frames_base = frames_per_chunk
+                else: # Last chunk takes all remaining frames
+                    processed_frames_base = (num_chunks - 1) * frames_per_chunk
+                    current_frames_base = total_frames - processed_frames_base
                 
+                current_frames = current_frames_base + (overlap_frames if i > 0 else 0)
+                # Ensure final frame count for generation is N*8+1
+                current_frames = self._align(current_frames, alignment_rule='n*8+1')
+
                 current_conditions = kwargs.get("initial_conditions", []) if i == 0 else []
                 if overlap_condition_item: current_conditions.append(overlap_condition_item)
 
@@ -231,7 +222,8 @@ class VideoService:
                     overlap_latents = chunk_latents[:, :, -overlap_frames:, :, :].clone()
                     overlap_condition_item = ConditioningItem(media_item=overlap_latents, media_frame_number=0, conditioning_strength=1.0)
                 
-                if i > 0: chunk_latents = chunk_latents[:, :, overlap_frames:, :, :]
+                if i > 0:
+                    chunk_latents = chunk_latents[:, :, overlap_frames:, :, :]
                 
                 chunk_path = RESULTS_DIR / f"temp_chunk_{i}_{used_seed}.pt"
                 torch.save(chunk_latents.cpu(), chunk_path)
@@ -359,12 +351,16 @@ class VideoService:
         else: self.runtime_autocast_dtype = torch.float32
         logging.info(f"Runtime precision policy set for autocast: {self.runtime_autocast_dtype}")
 
-    def _align(self, dim: int, alignment: int = FRAMES_ALIGNMENT) -> int:
+    def _align(self, dim: int, alignment: int = FRAMES_ALIGNMENT, alignment_rule: str = 'default') -> int:
+        """Aligns a dimension to the nearest multiple of `alignment`."""
+        if alignment_rule == 'n*8+1':
+             return ((dim - 1) // alignment) * alignment + 1
         return ((dim - 1) // alignment + 1) * alignment
     
     def _calculate_aligned_frames(self, duration_s: float, min_frames: int = 1) -> int:
         num_frames = int(round(duration_s * DEFAULT_FPS))
-        aligned_frames = self._align(num_frames)
+        # Para a duração total, sempre arredondamos para cima para o múltiplo de 8 mais próximo
+        aligned_frames = self._align(num_frames, alignment=FRAMES_ALIGNMENT)
         return max(aligned_frames, min_frames)
 
     def _get_random_seed(self) -> int: