Update deformes4D_engine.py

deformes4D_engine.py

@@ -20,6 +20,7 @@ from dataclasses import dataclass
 import gradio as gr
 import subprocess
 import gc
+import shutil
 
 from ltx_manager_helpers import ltx_manager_singleton
 from gemini_helpers import gemini_singleton 
@@ -49,6 +50,9 @@ class Deformes4DEngine:
         self._vae = None
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         logger.info("Especialista Deformes4D (Executor ADUC-SDR) inicializado.")
+        # Cria o diretório de workspace se não existir
+        os.makedirs(self.workspace_dir, exist_ok=True)
+
 
     @property
     def vae(self):
@@ -93,13 +97,22 @@ class Deformes4DEngine:
         list_file_path = os.path.join(self.workspace_dir, "concat_list.txt")
         with open(list_file_path, 'w', encoding='utf-8') as f:
             for path in video_paths: f.write(f"file '{os.path.abspath(path)}'\n")
-        cmd_list = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', list_file_path, '-c', 'copy', output_path]
+        
+        # Tenta usar aceleração de hardware (GPU) para a concatenação, se disponível
+        cmd_list = ['ffmpeg', '-y', '-hwaccel', 'auto', '-f', 'concat', '-safe', '0', '-i', list_file_path, '-c', 'copy', output_path]
         logger.info(f"Concatenando {len(video_paths)} clipes de vídeo em {output_path}...")
         try:
             subprocess.run(cmd_list, check=True, capture_output=True, text=True)
         except subprocess.CalledProcessError as e:
             logger.error(f"Erro no FFmpeg: {e.stderr}")
-            raise gr.Error(f"Falha na montagem final do vídeo. Detalhes: {e.stderr}")
+            # Tenta novamente sem aceleração de hardware como fallback
+            logger.info("Tentando concatenar novamente sem aceleração de hardware...")
+            cmd_list = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', list_file_path, '-c', 'copy', output_path]
+            try:
+                subprocess.run(cmd_list, check=True, capture_output=True, text=True)
+            except subprocess.CalledProcessError as e_fallback:
+                 logger.error(f"Erro no FFmpeg (fallback): {e_fallback.stderr}")
+                 raise gr.Error(f"Falha na montagem final do vídeo. Detalhes: {e_fallback.stderr}")
         
     # --- NÚCLEO DA LÓGICA ADUC-SDR ---
     def generate_full_movie(self, keyframes: list, global_prompt: str, storyboard: list, 
@@ -109,10 +122,18 @@ class Deformes4DEngine:
                             video_resolution: int, use_continuity_director: bool, 
                             progress: gr.Progress = gr.Progress()):
         
+        # --- ETAPA 0: SETUP ---
         FPS = 24
         FRAMES_PER_LATENT_CHUNK = 8
         ECO_LATENT_CHUNKS = 2
+        LATENT_PROCESSING_CHUNK_SIZE = 10 # Processa 10 fragmentos latentes por vez para economizar memória
         
+        run_timestamp = int(time.time())
+        temp_latent_dir = os.path.join(self.workspace_dir, f"temp_latents_{run_timestamp}")
+        temp_video_clips_dir = os.path.join(self.workspace_dir, f"temp_clips_{run_timestamp}")
+        os.makedirs(temp_latent_dir, exist_ok=True)
+        os.makedirs(temp_video_clips_dir, exist_ok=True)
+
         total_frames_brutos = self._quantize_to_multiple(int(seconds_per_fragment * FPS), FRAMES_PER_LATENT_CHUNK)
         frames_a_podar = self._quantize_to_multiple(int(total_frames_brutos * (trim_percent / 100)), FRAMES_PER_LATENT_CHUNK)
         latents_a_podar = frames_a_podar // FRAMES_PER_LATENT_CHUNK
@@ -128,17 +149,19 @@ class Deformes4DEngine:
         target_resolution_tuple = (video_resolution, video_resolution) 
         
         eco_latent_for_next_loop, dejavu_latent_for_next_loop = None, None
-        latent_fragments, latent_fragment_lengths = [], []
+        latent_fragment_paths = [] # Lista para armazenar caminhos dos latentes salvos no disco
         
         if len(keyframe_paths) < 2: raise gr.Error(f"A geração requer no mínimo 2 keyframes. Você forneceu {len(keyframe_paths)}.")
         
         num_transitions_to_generate = len(keyframe_paths) - 1
         
+        # --- ETAPA 1: GERAR FRAGMENTOS LATENTES E SALVAR EM DISCO ---
+        logger.info("--- INICIANDO ETAPA 1: Geração de Fragmentos Latentes ---")
         for i in range(num_transitions_to_generate):
             fragment_index = i + 1
-            progress(i / num_transitions_to_generate, desc=f"Gerando Latentes {fragment_index}/{num_transitions_to_generate}")
+            progress(i / num_transitions_to_generate, desc=f"Gerando Latente {fragment_index}/{num_transitions_to_generate}")
             
-            # ... (Lógica de decisão do Gemini e preparação de âncoras) ...
+            # (Lógica de decisão do Gemini e preparação de âncoras - inalterada)
             past_keyframe_path = keyframe_paths[i - 1] if i > 0 else keyframe_paths[i]
             start_keyframe_path = keyframe_paths[i]
             destination_keyframe_path = keyframe_paths[i + 1]
@@ -172,57 +195,110 @@ class Deformes4DEngine:
             if transition_type == "cut":
                 eco_latent_for_next_loop, dejavu_latent_for_next_loop = None, None
 
-            if use_upscaler:
-                # [REATORADO] Chamada para o novo especialista
-                latents_video = latent_enhancer_specialist_singleton.upscale(latents_video)
-            
-            latent_fragments.append(latents_video)
-            latent_fragment_lengths.append(latents_video.shape[2])
+            # [REATORADO] Mover latente para CPU e salvar no disco para liberar VRAM
+            cpu_latent = latents_video.cpu()
+            latent_path = os.path.join(temp_latent_dir, f"latent_fragment_{i:04d}.pt")
+            torch.save(cpu_latent, latent_path)
+            latent_fragment_paths.append(latent_path)
+
+            del latents_video, cpu_latent
+            gc.collect()
 
         del eco_latent_for_next_loop, dejavu_latent_for_next_loop
         gc.collect(); torch.cuda.empty_cache()
         
-        logger.info("--- CONCATENANDO E REFINANDO SUPER-LATENTE ---")
-        tensors_para_concatenar = [frag.to(self.device)[:, :, :-1, :, :] if i < len(latent_fragments) - 1 else frag.to(self.device) for i, frag in enumerate(latent_fragments)]
-        del latent_fragments; gc.collect(); torch.cuda.empty_cache()
+        # --- ETAPA 2: PROCESSAR LATENTES EM LOTES (CHUNKS) ---
+        logger.info(f"--- INICIANDO ETAPA 2: Processamento de {len(latent_fragment_paths)} latentes em lotes de {LATENT_PROCESSING_CHUNK_SIZE} ---")
+        final_video_clip_paths = []
+        num_chunks = -(-len(latent_fragment_paths) // LATENT_PROCESSING_CHUNK_SIZE) # Ceiling division
+
+        for i in range(num_chunks):
+            chunk_start_index = i * LATENT_PROCESSING_CHUNK_SIZE
+            chunk_end_index = chunk_start_index + LATENT_PROCESSING_CHUNK_SIZE
+            chunk_paths = latent_fragment_paths[chunk_start_index:chunk_end_index]
+            
+            progress(i / num_chunks, desc=f"Processando Lote {i+1}/{num_chunks}")
 
-        processed_latents = torch.cat(tensors_para_concatenar, dim=2)
-        del tensors_para_concatenar; gc.collect(); torch.cuda.empty_cache()
+            # Carrega os tensores do lote atual do disco para a GPU
+            tensors_in_chunk = [torch.load(p, map_location=self.device) for p in chunk_paths]
 
-        logger.info(f"Concatenação concluída. Shape do super-latente: {processed_latents.shape}")
+            # Concatena os tensores do lote, removendo o latente de sobreposição
+            tensors_para_concatenar = [
+                frag[:, :, :-1, :, :] if j < len(tensors_in_chunk) - 1 else frag
+                for j, frag in enumerate(tensors_in_chunk)
+            ]
+            sub_group_latent = torch.cat(tensors_para_concatenar, dim=2)
+            del tensors_in_chunk, tensors_para_concatenar
+            gc.collect(); torch.cuda.empty_cache()
 
-        base_name = f"fragment_{i}_{int(time.time())}"
-        current_path = os.path.join(self.workspace_dir, f"{base_name}_temp.mp4")
+            logger.info(f"Lote {i+1} concatenado. Shape do sub-latente: {sub_group_latent.shape}")
 
-        if use_audio:
-            current_path = self._generate_video_and_audio_from_latents(processed_latents, " ", base_name)
-        else:
-            current_path = self.latents_to_pixels(processed_latents)
-            self.save_video_from_tensor(pixel_tensor, current_path, fps=24)
-            del pixel_tensor
-                
+            # 1. (Opcional) Upscaler Latente
+            if use_upscaler:
+                logger.info(f"Aplicando Upscaler no lote {i+1}...")
+                sub_group_latent = latent_enhancer_specialist_singleton.upscale(sub_group_latent)
+                gc.collect(); torch.cuda.empty_cache()
+
+            # 2. Decodificar Latente para Vídeo (com ou sem áudio)
+            base_name = f"clip_{i:04d}_{run_timestamp}"
+            current_clip_path = os.path.join(temp_video_clips_dir, f"{base_name}_temp.mp4")
+            
+            if use_audio:
+                # O áudio é gerado para o prompt global por enquanto. Pode ser adaptado.
+                current_clip_path = self._generate_video_and_audio_from_latents(sub_group_latent, global_prompt, base_name)
+            else:
+                pixel_tensor = self.latents_to_pixels(sub_group_latent)
+                self.save_video_from_tensor(pixel_tensor, current_clip_path, fps=FPS)
+                del pixel_tensor
+            
+            del sub_group_latent
+            gc.collect(); torch.cuda.empty_cache()
+
+            # 3. (Opcional) Masterização HD
+            if use_hd:
+                logger.info(f"Aplicando masterização HD no clipe {i+1}...")
+                hd_clip_path = os.path.join(temp_video_clips_dir, f"{base_name}_hd.mp4")
+                try:
+                    hd_specialist_singleton.process_video(input_video_path=current_clip_path, output_video_path=hd_clip_path, prompt=global_prompt)
+                    # Apaga o clipe não-HD para economizar espaço
+                    if os.path.exists(current_clip_path) and current_clip_path != hd_clip_path:
+                        os.remove(current_clip_path)
+                    current_clip_path = hd_clip_path
+                except Exception as e:
+                    logger.error(f"Falha na masterização HD do clipe {i+1}: {e}. Usando versão padrão.")
+            
+            # 4. Adicionar caminho do clipe final à lista
+            final_video_clip_paths.append(current_clip_path)
 
-        if use_hd:
-            current_path = os.path.join(self.workspace_dir, f"{base_name}_hd.mp4")
-            try:
-                 hd_specialist_singleton.process_video(input_video_path=current_path, output_video_path=current_path, prompt=" ")
-            except Exception as e:
-                 logger.error(f"Falha na masterização HD do fragmento {i+1}: {e}. Usando versão padrão.")
-                        
         #if use_refiner:
         #    progress(0.8, desc="Refinando continuidade visual...")
         #    # [REATORADO] Chamada para o novo especialista
-        #
-        
-         
+        #    # OBS: Refinamento foi desativado conforme solicitado por degradar a lógica das keyframes.
+
+        # --- ETAPA 3: MONTAGEM FINAL ---
+        progress(0.98, desc="Montagem final dos clipes...")
+        final_video_path = os.path.join(self.workspace_dir, f"filme_final_{run_timestamp}.mp4")
+        self.concatenate_videos_ffmpeg(final_video_clip_paths, final_video_path)
         
-        progress(0.98, desc="Montagem final...")
-                   
-        logger.info(f"Processo concluído! Vídeo final salvo em: {current_path}")
-        yield {"final_path": current_path}
+        # --- ETAPA 4: LIMPEZA ---
+        logger.info("Limpando arquivos temporários...")
+        try:
+            shutil.rmtree(temp_latent_dir)
+            shutil.rmtree(temp_video_clips_dir)
+            concat_list_path = os.path.join(self.workspace_dir, "concat_list.txt")
+            if os.path.exists(concat_list_path):
+                os.remove(concat_list_path)
+        except OSError as e:
+            logger.warning(f"Não foi possível remover os diretórios temporários: {e}")
+
+        logger.info(f"Processo concluído! Vídeo final salvo em: {final_video_path}")
+        yield {"final_path": final_video_path}
 
     def _generate_video_and_audio_from_latents(self, latent_tensor, audio_prompt, base_name):
-        silent_video_path = os.path.join(self.workspace_dir, f"{base_name}_silent_for_audio.mp4")
+        # Este método agora opera em um diretório temporário para os clipes
+        temp_video_clips_dir = os.path.dirname(os.path.join(self.workspace_dir, base_name)) # Hack para obter o diretório correto
+        silent_video_path = os.path.join(temp_video_clips_dir, f"{base_name}_silent.mp4")
+
         pixel_tensor = self.latents_to_pixels(latent_tensor)
         self.save_video_from_tensor(pixel_tensor, silent_video_path, fps=24)
         del pixel_tensor; gc.collect(); torch.cuda.empty_cache()
@@ -233,13 +309,17 @@ class Deformes4DEngine:
                 capture_output=True, text=True, check=True)
             frag_duration = float(result.stdout.strip())
         except (subprocess.CalledProcessError, ValueError, FileNotFoundError):
-             logger.warning(f"ffprobe falhou. Calculando duração manualmente.")
-             num_pixel_frames = latent_tensor.shape[2] * 8
+             logger.warning(f"ffprobe falhou. Calculando duração manualmente a partir dos latentes.")
+             # O VAE interpola, então o número de frames é (num_latentes - 1) * 8 + 1 (aproximadamente)
+             num_pixel_frames = (latent_tensor.shape[2] - 1) * 8 + 1
              frag_duration = num_pixel_frames / 24.0
 
+        # Salva o vídeo com áudio no mesmo diretório temporário
         video_with_audio_path = audio_specialist_singleton.generate_audio_for_video(
             video_path=silent_video_path, prompt=audio_prompt,
-            duration_seconds=frag_duration)
+            duration_seconds=frag_duration,
+            output_path_override=os.path.join(temp_video_clips_dir, f"{base_name}_with_audio.mp4")
+        )
         
         if os.path.exists(silent_video_path):
              os.remove(silent_video_path)