Update api/ltx_server.py

api/ltx_server.py

@@ -397,89 +397,44 @@ class VideoService:
         return out
 
 
-    def _dividir_latentes_em_partes(self, latents_brutos, quantidade: int):
+    def _dividir_latentes_por_tamanho(self, latents_brutos, num_latente_por_chunk: int, overlap: int = 1):
         """
-        Divide um tensor de latentes em `quantidade` partes e retorna uma lista de clones.
-        Args: latents_brutos: tensor [B, C, T, H, W]
-              quantidade: número de partes que queremos dividir
-        Returns:
-            List[Tensor]: lista de `quantidade` partes, cada uma cloneada
-        """
-        sum_latent = latents_brutos.shape[2]  # dimensão temporal (latentes)
-        chunks = []
-
-        # minimo 5 latentes para dividir 
-        #min = 5
-        #if sum_latent < min or sum_latent//quantidade < min :
-        #    return [latents_brutos.clone()]
-    
-        # calcular tamanho aproximado de cada parte
-        step = sum_latent // quantidade
-        overlap = 1  
-        
-        print(f"================PODA CAUSAL=================")
-        print(f"[DEBUG] TOTAL LATENTES A DIVIDIR /{quantidade}  = {sum_latent}")
+        Divide o tensor de latentes em chunks com tamanho definido em número de latentes.
         
-        for i in range(quantidade):
-            start = i * step 
-            end = start + step
-            if i == quantidade - 1:
-                end = sum_latent
-            else:
-                end += overlap  
-            if i > 0:
-                start = start - overlap
-            chunk = latents_brutos[:, :, start:end, :, :].clone()
-            chunks.append(chunk)
-            chunk_sum = chunk.shape[2]  
-            print(f"[DEBUG] chunk{i+1}[:, :, {start}:{end}, :, :]  = {chunk_sum}")
-
-        print(f"================PODA CAUSAL=================")
-        return chunks
-    
-
-    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
-    
+            num_latente_por_chunk: número de latentes por chunk
+            overlap: número de frames que se sobrepõem entre chunks
+        
         Returns:
-            List[Tensor]: lista de partes do tensor (cada uma cloneada)
+            List[tensor]: lista de chunks cloneados
         """
-        total_latents = latents_brutos.shape[2]
-        overlap = 1
+        sum_latent = latents_brutos.shape[2]
         chunks = []
     
-        print("================PODA CAUSAL=================")
-        print(f"[DEBUG] TOTAL LATENTES A DIVIDIR  = {total_latents}")
-        print(f"[DEBUG] Num LATENTES por chunk  = {min_latentes_por_chunk}")
+        if num_latente_por_chunk >= sum_latent:
+            return [latents_brutos.clone()]
     
-        start = 0
-        i = 0
+        steps = (sum_latent + num_latente_por_chunk - 1) // num_latente_por_chunk  # ceil
+        print("================PODA CAUSAL=================")
+        print(f"[DEBUG] TOTAL LATENTES = {sum_latent}")
+        print(f"[DEBUG] Num LATENTES por chunk = {num_latente_por_chunk}")
+        print(f"[DEBUG] Número de chunks = {steps}")
     
-        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
+        for i in range(steps):
+            start = i * num_latente_por_chunk
+            end = start + num_latente_por_chunk
+            if i > 0:
+                start -= overlap  # sobreposição
+            if end > sum_latent:
+                end = sum_latent
             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(f"[DEBUG] chunk{i+1}[:, :, {start}:{end}, :, :] = {chunk.shape[2]}")
     
         print("================PODA CAUSAL=================")
         return chunks
+
         
     def _concat_crossfade_videos_por_frames(self, mp4_list, out_path, fps=24, crossfade_frames=8):
         """
@@ -749,7 +704,7 @@ class VideoService:
             #lat_a1, lat_a2 = self._dividir_latentes(lat_a)
             #lat_b1, lat_b2 = self._dividir_latentes(lat_b)
             #latents_parts = [lat_a1, lat_a2, lat_b1, lat_b2]
-            latents_parts = self._dividir_latentes_em_partes(latents_cpu,4)
+            latents_parts = self._dividir_latentes_por_tamanho(latents_cpu,6,1)
             
             temp_dir = tempfile.mkdtemp(prefix="ltxv_"); self._register_tmp_dir(temp_dir)
             results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)