Update aduc_framework/engineers/deformes4D.py

aduc_framework/engineers/deformes4D.py

@@ -1,190 +1,235 @@
-# aduc_framework/orchestrator.py
+# aduc_framework/engineers/deformes4D.py
 #
 # Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
 #
-# Versão 3.2.1 (Com lógica de Upscale por fragmento)
+# Versão 3.1.1 (Com correção de limpeza de arquivos)
 #
-# Esta versão representa a camada de orquestração do Aduc Framework.
-# Ela é agnóstica a qualquer interface (UI ou API) e opera com
-# tipos de dados bem definidos (Pydantic) e um estado de geração central.
+# Este engenheiro implementa a Câmera (Ψ) e o Destilador (Δ) da arquitetura
+# ADUC-SDR. Sua única responsabilidade é a geração sequencial de fragmentos de
+# vídeo com base em um conjunto de keyframes pré-definido.
 
+import os
+import time
+import imageio
+import numpy as np
+import torch
 import logging
-from typing import List, Dict, Any, Tuple, Callable, Optional, Generator
 from PIL import Image, ImageOps
-import os
-import subprocess
+import gc
 import shutil
 from pathlib import Path
-import time
-import gc
-import torch
+from typing import List, Tuple, Dict, Any, Callable, Optional
 
-# Importa componentes internos do framework
-from .director import AducDirector
-from .types import GenerationState, PreProductionParams, ProductionParams
-from .engineers import deformes2d_thinker_singleton, deformes3d_engine_singleton, Deformes4DEngine
-# Importa managers diretamente para as tarefas de pós-produção
-from .managers.latent_enhancer_manager import latent_enhancer_specialist_singleton
-from .managers.seedvr_manager import seedvr_manager_singleton
-from .managers.mmaudio_manager import mmaudio_manager_singleton
-from .managers.vae_manager import vae_manager_singleton
-from .tools.video_encode_tool import video_encode_tool_singleton
+# --- Imports Relativos Corrigidos ---
+from ..types import LatentConditioningItem
+from ..managers.ltx_manager import ltx_manager_singleton
+from ..managers.vae_manager import vae_manager_singleton
+from .deformes2D_thinker import deformes2d_thinker_singleton
+from ..tools.video_encode_tool import video_encode_tool_singleton
 
 logger = logging.getLogger(__name__)
 
 ProgressCallback = Optional[Callable[[float, str], None]]
 
-class AducOrchestrator:
+class Deformes4DEngine:
     """
-    Implementa o Maestro (Γ), a camada de orquestração central do Aduc Framework.
-    Ele recebe solicitações, atualiza o estado de geração, delega tarefas para os
-    engenheiros especialistas e retorna o estado atualizado.
+    Orquestra a geração e concatenação de fragmentos de vídeo.
     """
-    def __init__(self, workspace_dir: str):
-        self.director = AducDirector(workspace_dir)
-        self.editor = Deformes4DEngine()
-        self.editor.initialize(workspace_dir)
-        self.painter = deformes3d_engine_singleton
-        self.painter.initialize(workspace_dir)
+    def __init__(self):
+        """O construtor é leve e não recebe argumentos."""
+        self.workspace_dir: Optional[str] = None
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        logger.info("ADUC Maestro (Framework Core) está no pódio. Engenheiros especialistas prontos.")
-
-    def get_current_state(self) -> GenerationState:
-        return self.director.get_full_state()
-
-    def process_image_for_story(self, image_path: str, size: int, filename: str) -> str:
-        img = Image.open(image_path).convert("RGB")
-        img_square = ImageOps.fit(img, (size, size), Image.Resampling.LANCZOS)
-        processed_path = os.path.join(self.director.workspace_dir, filename)
-        img_square.save(processed_path)
-        logger.info(f"Imagem de referência processada e salva em: {processed_path}")
-        return processed_path
-
-    def task_pre_production(self, params: PreProductionParams, progress_callback: ProgressCallback = None) -> Tuple[List[str], List[str], GenerationState]:
-        logger.info("Maestro: Iniciando tarefa de Pré-Produção.")
-        self.director.update_parameters("pre_producao", params)
-        if progress_callback: progress_callback(0.1, "Gerando storyboard...")
-        storyboard_list = deformes2d_thinker_singleton.generate_storyboard(prompt=params.prompt, num_keyframes=params.num_keyframes, ref_image_paths=params.ref_paths)
-        self.director.update_pre_production_state(params.prompt, params.ref_paths, storyboard_list)
-        if progress_callback: progress_callback(0.2, "Iniciando geração de keyframes...")
-        keyframes_detailed_data = self.painter.generate_keyframes_from_storyboard(generation_state=self.director.get_full_state_as_dict(), progress_callback=progress_callback)
-        self.director.update_keyframes_state(keyframes_detailed_data)
-        final_keyframe_paths = [kf["caminho_pixel"] for kf in keyframes_detailed_data]
-        final_state = self.director.get_full_state()
-        logger.info("Maestro: Tarefa de Pré-Produção concluída.")
-        return storyboard_list, final_keyframe_paths, final_state
-
-    def task_produce_original_movie(self, params: ProductionParams, progress_callback: ProgressCallback = None) -> Tuple[str, List[str], GenerationState]:
-        logger.info("Maestro: Iniciando tarefa de Produção do Filme Original.")
-        self.director.update_parameters("producao", params)
-        result_data = self.editor.generate_original_movie(full_generation_state=self.director.get_full_state_as_dict(), progress_callback=progress_callback)
-        self.director.update_video_state(result_data["video_data"])
-        final_video_path = result_data["final_path"]
-        latent_paths = result_data["latent_paths"]
-        final_state = self.director.get_full_state()
-        logger.info("Maestro: Tarefa de Produção do Filme Original concluída.")
-        return final_video_path, latent_paths, final_state
-
-    # --- TAREFAS DE PÓS-PRODUÇÃO ---
-
-    def task_run_latent_upscaler(self, latent_paths: List[str], chunk_size: int, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
-        if not latent_paths:
-            raise ValueError("Não é possível fazer o upscale: nenhum caminho de latente foi fornecido.")
+        logger.info("Deformes4DEngine instanciado (não inicializado).")
+
+    def initialize(self, workspace_dir: str):
+        """Inicializa o engenheiro com as configurações necessárias."""
+        if self.workspace_dir is not None:
+            return # Evita reinicialização
+        self.workspace_dir = workspace_dir
+        os.makedirs(self.workspace_dir, exist_ok=True)
+        logger.info(f"Deformes4D Specialist (Executor) inicializado com workspace: {self.workspace_dir}.")
+
+    def generate_original_movie(
+        self,
+        full_generation_state: Dict[str, Any],
+        progress_callback: ProgressCallback = None
+    ) -> Dict[str, Any]:
+        """
+        Gera o filme principal lendo todos os parâmetros do estado de geração.
+        """
+        if not self.workspace_dir:
+            raise RuntimeError("Deformes4DEngine não foi inicializado. Chame o método initialize() antes de usar.")
+
+        # 1. Extrai todos os parâmetros do estado de geração
+        pre_prod_params = full_generation_state.get("parametros_geracao", {}).get("pre_producao", {})
+        prod_params = full_generation_state.get("parametros_geracao", {}).get("producao", {})
         
-        logger.info("--- ORQUESTRADOR: Tarefa de Upscaling de Latentes (Estratégia por Fragmento) ---")
-        run_timestamp = int(time.time())
-        temp_dir = os.path.join(self.director.workspace_dir, f"temp_upscaled_clips_{run_timestamp}")
-        os.makedirs(temp_dir, exist_ok=True)
+        keyframes_data = full_generation_state.get("Keyframe_atos", [])
+        global_prompt = full_generation_state.get("Promt_geral", "")
+        storyboard = [ato["resumo_ato"] for ato in full_generation_state.get("Atos", [])]
+        keyframe_paths = [kf["caminho_pixel"] for kf in keyframes_data]
         
-        final_upscaled_clip_paths = []
-        total_latents = len(latent_paths)
+        seconds_per_fragment = pre_prod_params.get('duration_per_fragment', 4.0)
+        video_resolution = pre_prod_params.get('resolution', 480)
         
-        for i, latent_path in enumerate(latent_paths):
-            if progress_callback:
-                progress_fraction = i / total_latents
-                progress_callback(progress_fraction, f"Processando fragmento latente {i+1}/{total_latents}")
-
-            try:
-                logger.info(f"Carregando latente: {os.path.basename(latent_path)}")
-                latent_tensor = torch.load(latent_path, map_location=self.device)
-                
-                logger.info(f"Fazendo upscale do latente shape: {latent_tensor.shape}")
-                upscaled_latent = latent_enhancer_specialist_singleton.upscale(latent_tensor)
-                del latent_tensor; gc.collect(); torch.cuda.empty_cache()
-                
-                logger.info(f"Decodificando latente com upscale shape: {upscaled_latent.shape}")
-                pixel_tensor = vae_manager_singleton.decode(upscaled_latent)
-                del upscaled_latent; gc.collect(); torch.cuda.empty_cache()
-
-                clip_filename = f"upscaled_clip_{i:04d}.mp4"
-                current_clip_path = os.path.join(temp_dir, clip_filename)
-                self.editor.save_video_from_tensor(pixel_tensor, current_clip_path, fps=24)
-                final_upscaled_clip_paths.append(current_clip_path)
-                logger.info(f"Clipe com upscale salvo em: {clip_filename}")
-                del pixel_tensor; gc.collect(); torch.cuda.empty_cache()
-
-                yield {"progress": (i + 1) / total_latents, "desc": f"Fragmento {i+1}/{total_latents} concluído."}
-
-            except Exception as e:
-                logger.error(f"Falha ao processar o fragmento latente {latent_path}: {e}", exc_info=True)
-                continue
-
-        if progress_callback: progress_callback(0.98, "Montando vídeo final com upscale...")
-        final_video_path = os.path.join(self.director.workspace_dir, f"upscaled_movie_{run_timestamp}.mp4")
-        video_encode_tool_singleton.concatenate_videos(final_upscaled_clip_paths, final_video_path, self.director.workspace_dir)
+        trim_percent = prod_params.get('trim_percent', 50)
+        handler_strength = prod_params.get('handler_strength', 0.5)
+        destination_convergence_strength = prod_params.get('destination_convergence_strength', 0.75)
+        guidance_scale = prod_params.get('guidance_scale', 2.0)
+        stg_scale = prod_params.get('stg_scale', 0.025)
+        num_inference_steps = prod_params.get('inference_steps', 20)
         
-        try:
-            shutil.rmtree(temp_dir)
-        except OSError as e:
-            logger.warning(f"Não foi possível remover o diretório temporário de clipes com upscale: {e}")
-            
-        logger.info(f"Upscaling de latentes completo! Vídeo final em: {final_video_path}")
-        yield {"final_path": final_video_path}
+        # 2. Inicia o processo de geração
+        FPS = 24
+        FRAMES_PER_LATENT_CHUNK = 8
+        LATENT_PROCESSING_CHUNK_SIZE = 4
 
-    def task_run_hd_mastering(self, source_video_path: str, model_version: str, steps: int, prompt: str, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
-        logger.info(f"--- ORQUESTRADOR: Tarefa de Masterização HD com SeedVR {model_version} ---")
-        
         run_timestamp = int(time.time())
-        output_path = os.path.join(self.director.workspace_dir, f"hd_mastered_movie_{model_version}_{run_timestamp}.mp4")
-        
-        final_path = seedvr_manager_singleton.process_video(
-            input_video_path=source_video_path,
-            output_video_path=output_path,
-            prompt=prompt,
-            model_version=model_version,
-            steps=steps,
-            progress=progress_callback
-        )
-        logger.info(f"Masterização HD completa! Vídeo final em: {final_path}")
-        yield {"final_path": final_path}
-
-    def task_run_audio_generation(self, source_video_path: str, audio_prompt: str, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
-        logger.info(f"--- ORQUESTRADOR: Tarefa de Geração de Áudio ---")
-
-        if progress_callback: progress_callback(0.1, "Preparando para geração de áudio...")
-        
-        run_timestamp = int(time.time())
-        source_name = Path(source_video_path).stem
-        output_path = os.path.join(self.director.workspace_dir, f"{source_name}_with_audio_{run_timestamp}.mp4")
+        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
+        DEJAVU_FRAME_TARGET = frames_a_podar - 1 if frames_a_podar > 0 else 0
+        DESTINATION_FRAME_TARGET = total_frames_brutos - 1
+
+        base_ltx_params = {"guidance_scale": guidance_scale, "stg_scale": stg_scale, "num_inference_steps": num_inference_steps}
+        story_history = ""
+        target_resolution_tuple = (video_resolution, video_resolution)
+        eco_latent_for_next_loop, dejavu_latent_for_next_loop = None, None
+        latent_fragment_paths = []
+        video_fragments_data = []
+
+        if len(keyframe_paths) < 2:
+            raise ValueError(f"A geração requer pelo menos 2 keyframes. Fornecidos: {len(keyframe_paths)}.")
+        num_transitions_to_generate = len(keyframe_paths) - 1
+
+        logger.info("--- ESTÁGIO 1: Geração de Fragmentos Latentes ---")
+        for i in range(num_transitions_to_generate):
+            fragment_index = i + 1
+            if progress_callback:
+                progress_fraction = (i / num_transitions_to_generate) * 0.7
+                progress_callback(progress_fraction, f"Gerando Latente {fragment_index}/{num_transitions_to_generate}")
+
+            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]
+            future_story_prompt = storyboard[i + 1] if (i + 1) < len(storyboard) else "A cena final."
+            decision = deformes2d_thinker_singleton.get_cinematic_decision(
+                global_prompt, story_history, past_keyframe_path, start_keyframe_path, 
+                destination_keyframe_path, storyboard[i - 1] if i > 0 else "O início.", 
+                storyboard[i], future_story_prompt
+            )
+            motion_prompt = decision["motion_prompt"]
+            story_history += f"\n- Ato {fragment_index}: {motion_prompt}"
+            
+            conditioning_items = []
+            if eco_latent_for_next_loop is None:
+               img_start = self._preprocess_image_for_latent_conversion(Image.open(start_keyframe_path).convert("RGB"), target_resolution_tuple)
+               conditioning_items.append(LatentConditioningItem(self._pil_to_latent(img_start), 0, 1.0))
+            else:
+               conditioning_items.append(LatentConditioningItem(eco_latent_for_next_loop, 0, 1.0))
+               conditioning_items.append(LatentConditioningItem(dejavu_latent_for_next_loop, DEJAVU_FRAME_TARGET, handler_strength))
+            
+            img_dest = self._preprocess_image_for_latent_conversion(Image.open(destination_keyframe_path).convert("RGB"), target_resolution_tuple)
+            conditioning_items.append(LatentConditioningItem(self._pil_to_latent(img_dest), DESTINATION_FRAME_TARGET, destination_convergence_strength))
+            
+            latents_brutos, _ = ltx_manager_singleton.generate_latent_fragment(
+                height=video_resolution, width=video_resolution,
+                conditioning_items_data=conditioning_items, motion_prompt=motion_prompt,
+                video_total_frames=total_frames_brutos, video_fps=FPS,
+                **base_ltx_params
+            )
+            
+            last_trim = latents_brutos[:, :, -(latents_a_podar+1):, :, :].clone()
+            eco_latent_for_next_loop = last_trim[:, :, :2, :, :].clone()
+            dejavu_latent_for_next_loop = last_trim[:, :, -1:, :, :].clone()
+            latents_video = latents_brutos[:, :, :-(latents_a_podar-1), :, :].clone()
+            del last_trim, latents_brutos; gc.collect(); torch.cuda.empty_cache()
+            
+            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)
+            
+            video_fragments_data.append({"id": i, "prompt_video": motion_prompt})
+            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(f"--- ESTÁGIO 2: Processando {len(latent_fragment_paths)} latentes ---")
+        final_video_clip_paths = []
+        num_chunks = -(-len(latent_fragment_paths) // LATENT_PROCESSING_CHUNK_SIZE) if LATENT_PROCESSING_CHUNK_SIZE > 0 else 0
+        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]
+            
+            if progress_callback:
+                progress_fraction = 0.7 + (i / num_chunks * 0.28)
+                progress_callback(progress_fraction, f"Processando & Decodificando Lote {i+1}/{num_chunks}")
 
-        result = subprocess.run(
-            ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", source_video_path],
-            capture_output=True, text=True, check=True
-        )
-        duration = float(result.stdout.strip())
+            tensors_in_chunk = [torch.load(p, map_location=self.device) for p in chunk_paths]
+            sub_group_latent = torch.cat(tensors_in_chunk, dim=2)
+            del tensors_in_chunk; gc.collect(); torch.cuda.empty_cache()
 
-        if progress_callback: progress_callback(0.5, "Gerando trilha de áudio...")
+            pixel_tensor = vae_manager_singleton.decode(sub_group_latent)
+            del sub_group_latent; gc.collect(); torch.cuda.empty_cache()
+
+            base_name = f"clip_{i:04d}_{run_timestamp}"
+            current_clip_path = os.path.join(temp_video_clips_dir, f"{base_name}.mp4")
+            self.save_video_from_tensor(pixel_tensor, current_clip_path, fps=FPS)
+            final_video_clip_paths.append(current_clip_path)
+            del pixel_tensor; gc.collect(); torch.cuda.empty_cache()
+        
+        if progress_callback: progress_callback(0.98, "Montando o filme final...")
+        final_video_path = os.path.join(self.workspace_dir, f"original_movie_{run_timestamp}.mp4")
+        video_encode_tool_singleton.concatenate_videos(final_video_clip_paths, final_video_path, self.workspace_dir)
         
-        final_path = mmaudio_manager_singleton.generate_audio_for_video(
-            video_path=source_video_path,
-            prompt=audio_prompt,
-            duration_seconds=duration,
-            output_path_override=output_path
-        )
+        try:
+            shutil.rmtree(temp_video_clips_dir)
+            # A linha que apagava 'temp_latent_dir' foi removida para persistir os latentes.
+        except OSError as e:
+            logger.warning(f"Não foi possível remover o diretório de clipes temporários: {e}")
+            
+        logger.info(f"Processo completo! Vídeo original salvo em: {final_video_path}")
         
-        logger.info(f"Geração de áudio completa! Vídeo com áudio em: {final_path}")
-        if progress_callback: progress_callback(1.0, "Geração de áudio completa!")
-        yield {"final_path": final_path}
+        final_video_data_for_state = {
+            "id": 0, "caminho_pixel": final_video_path,
+            "caminhos_latentes_fragmentos": latent_fragment_paths,
+            "fragmentos_componentes": video_fragments_data
+        }
+
+        return {
+            "final_path": final_video_path,
+            "latent_paths": latent_fragment_paths,
+            "video_data": final_video_data_for_state
+        }
+
+    # --- FUNÇÕES HELPER ---
+    def save_video_from_tensor(self, video_tensor: torch.Tensor, path: str, fps: int = 24):
+        if video_tensor is None or video_tensor.ndim != 5 or video_tensor.shape[2] == 0: return
+        video_tensor = video_tensor.squeeze(0).permute(1, 2, 3, 0)
+        video_tensor = (video_tensor.clamp(-1, 1) + 1) / 2.0
+        video_np = (video_tensor.detach().cpu().float().numpy() * 255).astype(np.uint8)
+        with imageio.get_writer(path, fps=fps, codec='libx264', quality=8, output_params=['-pix_fmt', 'yuv420p']) as writer:
+            for frame in video_np: writer.append_data(frame)
+
+    def _preprocess_image_for_latent_conversion(self, image: Image.Image, target_resolution: tuple) -> Image.Image:
+        if image.size != target_resolution:
+            return ImageOps.fit(image, target_resolution, Image.Resampling.LANCZOS)
+        return image
+
+    def _pil_to_latent(self, pil_image: Image.Image) -> torch.Tensor:
+        image_np = np.array(pil_image).astype(np.float32) / 255.0
+        tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0).unsqueeze(2)
+        tensor = (tensor * 2.0) - 1.0
+        return vae_manager_singleton.encode(tensor)
+
+    def _quantize_to_multiple(self, n: int, m: int) -> int:
+        if m == 0: return n
+        quantized = int(round(n / m) * m)
+        return m if n > 0 and quantized == 0 else quantized