Update api/ltx_server.py

api/ltx_server.py

@@ -437,22 +437,49 @@ class VideoService:
         return chunks
     
 
-    def dividir_latentes(latents_brutos):
-        sum_latent = latents_brutos.shape[2]  # dimensão temporal (latentes)
-        cut1 = sum_latent // 2
-        cut2 = sum_latent-cut1-1
-        chunk1 = latents_brutos[:, :, :cut1, :, :].clone()
-        chunk2 = latents_brutos[:, :, cut2:, :, :].clone()
-        chunk1_sum = primeir.shape[2]  # dimensão temporal (latentes1)
-        chunk2_sum = segunda.shape[2]  # dimensão temporal (latentes2)
-     
-        print(f"==============PODA CAUSAL==================")
-        print(f"[DEBUG] TOTAL LATENTES A DIVIDIR    = {sum_latent}")
-        print(f"[DEBUG] chunk1[:, :, :{cut1}, :, :] = {chunk1_sum}")
-        print(f"[DEBUG] chunk2[:, :, :{cut2}, :, :] = {chunk2_sum}")
-        print(f"===========================================")
-        
-        return primeira, segunda
+    def _dividir_latentes_em_partes(self, latents_brutos, min_latentes_por_chunk: int):
+        """
+        Divide o tensor de latentes em várias partes com base no número mínimo de latentes por chunk.
+        Cada parte tem sobreposição causal de 1 frame com a anterior.
+    
+        Args:
+            latents_brutos: tensor [B, C, T, H, W]
+            min_latentes_por_chunk: número mínimo de latentes em cada parte
+    
+        Returns:
+            List[Tensor]: lista de partes do tensor (cada uma cloneada)
+        """
+        total_latents = latents_brutos.shape[2]
+        overlap = 1
+        chunks = []
+    
+        print("================PODA CAUSAL=================")
+        print(f"[DEBUG] TOTAL LATENTES A DIVIDIR  = {total_latents}")
+        print(f"[DEBUG] Num LATENTES por chunk  = {min_latentes_por_chunk}")
+    
+        start = 0
+        i = 0
+    
+        while start < total_latents:
+            end = start + min_latentes_por_chunk - 1  # menos 1 por causa da sobreposição
+            if end > total_latents:
+                end = total_latents
+            chunk = latents_brutos[:, :, start:end, :, :].clone()
+            chunks.append(chunk)
+    
+            print(f"[DEBUG] chunk{i+1}[:, :, {start}:{end}, :, :]  = {chunk.shape[2]}")
+    
+            # Avança com sobreposição
+            start = end - overlap
+            i += 1
+    
+            # Evita loop infinito se o final for atingido
+            if start >= total_latents:
+                break
+    
+        print("================PODA CAUSAL=================")
+        return chunks
+
 
         
     def _concat_mp4s_no_reencode(self, mp4_list: List[str], out_path: str):