deformes4D_engine.py

# deformes4D_engine.py
# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
#
# MODIFICATIONS FOR ADUC-SDR:
# Copyright (C) 2025 Carlos Rodrigues dos Santos. All rights reserved.
#
# This file is part of the ADUC-SDR project. It contains the core logic for
# video fragment generation, latent manipulation, and dynamic editing, 
# governed by the ADUC orchestrator.
# This component is licensed under the GNU Affero General Public License v3.0.

import os
import time
import imageio
import numpy as np
import torch
import logging
from PIL import Image, ImageOps
from dataclasses import dataclass
import gradio as gr
import subprocess
import gc


from audio_specialist import audio_specialist_singleton
from ltx_manager_helpers import ltx_manager_singleton
from gemini_helpers import gemini_singleton 
from upscaler_specialist import upscaler_specialist_singleton
from hd_specialist import hd_specialist_singleton # Importa o novo especialista
from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode

logger = logging.getLogger(__name__)

@dataclass
class LatentConditioningItem:
    """Representa uma âncora de condicionamento no espaço latente para a Câmera (Ψ)."""
    latent_tensor: torch.Tensor
    media_frame_number: int
    conditioning_strength: float

class Deformes4DEngine:
    """
    Implementa a Câmera (Ψ) e o Destilador (Δ) da arquitetura ADUC-SDR.
    Orquestra a geração, pós-produção latente e renderização final dos fragmentos de vídeo.
    """
    def __init__(self, ltx_manager, workspace_dir="deformes_workspace"):
        self.ltx_manager = ltx_manager
        self.workspace_dir = workspace_dir
        self._vae = None
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        logger.info("Especialista Deformes4D (Executor ADUC-SDR: Câmera Ψ e Destilador Δ) inicializado.")

    @property
    def vae(self):
        if self._vae is None:
            self._vae = self.ltx_manager.workers[0].pipeline.vae
        self._vae.to(self.device); self._vae.eval()
        return self._vae

    # MÉTODOS AUXILIARES
    def save_latent_tensor(self, tensor: torch.Tensor, path: str):
        torch.save(tensor.cpu(), path)

    def load_latent_tensor(self, path: str) -> torch.Tensor:
        return torch.load(path, map_location=self.device)

    @torch.no_grad()
    def pixels_to_latents(self, tensor: torch.Tensor) -> torch.Tensor:
        tensor = tensor.to(self.device, dtype=self.vae.dtype)
        return vae_encode(tensor, self.vae, vae_per_channel_normalize=True)

    @torch.no_grad()
    def latents_to_pixels(self, latent_tensor: torch.Tensor, decode_timestep: float = 0.05) -> torch.Tensor:
        latent_tensor = latent_tensor.to(self.device, dtype=self.vae.dtype)
        timestep_tensor = torch.tensor([decode_timestep] * latent_tensor.shape[0], device=self.device, dtype=latent_tensor.dtype)
        return vae_decode(latent_tensor, self.vae, is_video=True, timestep=timestep_tensor, vae_per_channel_normalize=True)

    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 self.pixels_to_latents(tensor)
    
    def _get_video_duration(self, video_path: str) -> float:
        if not os.path.exists(video_path): return 0.0
        try:
            result = subprocess.run(
                ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", video_path],
                capture_output=True, text=True, check=True)
            return float(result.stdout.strip())
        except Exception:
            return 0.0

    def _combine_video_and_audio_ffmpeg(self, video_path: str, audio_path: str, output_path: str):
        """Combina um arquivo de vídeo com um arquivo de áudio usando ffmpeg."""
        cmd = [
            'ffmpeg', '-y',
            '-i', video_path,
            '-i', audio_path,
            '-c:v', 'copy',      # Copia o stream de vídeo sem re-codificar
            '-c:a', 'aac',       # Re-codifica o áudio para o formato AAC, padrão para MP4
            '-shortest',         # Termina a codificação quando o stream mais curto terminar
            output_path
        ]
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True, encoding='utf-8')
            logger.info(f"Áudio e vídeo combinados com sucesso em {output_path}")
        except subprocess.CalledProcessError as e:
            logger.error(f"Falha ao combinar áudio e vídeo. Detalhes: {e.stderr}")
            raise gr.Error(f"Falha ao combinar áudio e vídeo: {e.stderr}")

    def _generate_standalone_audio(self, video_for_duration_path: str, audio_prompt: str) -> str:
        """Gera um arquivo de áudio e retorna seu caminho."""
        duration = self._get_video_duration(video_for_duration_path)
        if duration == 0:
            raise gr.Error("Não foi possível determinar a duração do vídeo para gerar o áudio.")
        
        # Esta função agora deve retornar apenas o caminho do arquivo de áudio gerado
        # (pode exigir uma pequena modificação no seu audio_specialist)
        audio_path = audio_specialist_singleton.generate_audio(
            prompt=audio_prompt,
            duration_seconds=duration,
            output_dir=self.workspace_dir
        )
        return audio_path

    # NÚCLEO DA LÓGICA ADUC-SDR
    def generate_full_movie(self, keyframes: list, global_prompt: str, storyboard: list, 
                            seconds_per_fragment: float, trim_percent: int,
                            handler_strength: float, destination_convergence_strength: float, 
                            video_resolution: int, use_continuity_director: bool, 
                            progress: gr.Progress = gr.Progress()):
        
        TOTAL_STEPS = len(keyframes) - 1 + 5 # Fragmentos + 5 etapas de pós-produção
        current_step = 0

        FPS = 24
        FRAMES_PER_LATENT_CHUNK = 8
        ECO_LATENT_CHUNKS = 2
        
        total_frames_brutos = self._quantize_to_multiple(int(seconds_per_fragment * FPS), FRAMES_PER_LATENT_CHUNK)
        total_latents_brutos = total_frames_brutos // 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

        if total_latents_brutos <= latents_a_podar + 1:
            raise gr.Error(f"A combinação de duração e poda é muito agressiva.")

        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": 2.0, "stg_scale": 0.025, "rescaling_scale": 0.15, "num_inference_steps": 20}
        keyframe_paths = [item[0] if isinstance(item, tuple) else item for item in keyframes]
        story_history = ""
        
        eco_latent_for_next_loop = None
        dejavu_latent_for_next_loop = None
        
        num_transitions_to_generate = len(keyframe_paths) - 1
        processed_latent_fragments = []

        for i in range(num_transitions_to_generate):
            fragment_index = i + 1
            current_step += 1
            progress(current_step / TOTAL_STEPS, desc=f"Gerando Fragmento {fragment_index}/{num_transitions_to_generate}")
            
            # ... (Lógica de decisão do Gemini e configuração de parâmetros - sem alterações)
            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 = gemini_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)
            transition_type, motion_prompt = decision["transition_type"], decision["motion_prompt"]
            story_history += f"\n- Ato {fragment_index}: {motion_prompt}"

            expected_height, expected_width = 768, 1152
            downscale_factor = 2 / 3
            downscaled_height = self._quantize_to_multiple(int(expected_height * downscale_factor), 8)
            downscaled_width = self._quantize_to_multiple(int(expected_width * downscale_factor), 8)
            target_resolution_tuple = (downscaled_height, downscaled_width)

            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))

            current_ltx_params = {**base_ltx_params, "motion_prompt": motion_prompt}
            latents_brutos, _ = self._generate_latent_tensor_internal(conditioning_items, current_ltx_params, target_resolution_tuple, total_frames_brutos)

            last_trim = latents_brutos[:, :, -(latents_a_podar+1):, :, :].clone()
            eco_latent_for_next_loop = last_trim[:, :, :ECO_LATENT_CHUNKS, :, :].clone()   
            dejavu_latent_for_next_loop = last_trim[:, :, -1:, :, :].clone()
            latents_video = latents_brutos[:, :, :-(latents_a_podar-1), :, :].clone()
            latents_video = latents_video[:, :, 1:, :, :]

            if transition_type == "cut":
                eco_latent_for_next_loop, dejavu_latent_for_next_loop = None, None

            # --- ATO I: PÓS-PRODUÇÃO LATENTE ---
            upscaled_latents = self.upscale_latents(latents_video)
            refined_latents = self.refine_latents(upscaled_latents, motion_prompt=f"refining scene: {motion_prompt}")
            processed_latent_fragments.append(refined_latents)

        # --- FIM DO LOOP DE GERAÇÃO ---

        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Concatenando fragmentos...")
        tensors_para_concatenar = [frag.to(self.device) for frag in processed_latent_fragments]
        final_concatenated_latents = torch.cat(tensors_para_concatenar, dim=2)

        base_name = f"movie_{int(time.time())}"
        
        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Renderizando vídeo base...")
        refined_silent_video_path = os.path.join(self.workspace_dir, f"{base_name}_refined_silent.mp4")
        final_pixel_tensor = self.latents_to_pixels(final_concatenated_latents)
        self.save_video_from_tensor(final_pixel_tensor, refined_silent_video_path, fps=FPS)
        
        # Limpeza de VRAM antes da próxima etapa pesada
        del final_pixel_tensor, final_concatenated_latents, processed_latent_fragments, tensors_para_concatenar
        gc.collect()
        torch.cuda.empty_cache()

        
        # --- ATO II: MASTERIZAÇÃO FINAL (APLICAÇÃO DE HD) ---
        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Aprimoramento final (HD)...")
        hq_silent_video_path = os.path.join(self.workspace_dir, f"{base_name}_hq_silent.mp4")
        try:
            # O Especialista HD processa o vídeo silencioso refinado
            hd_specialist_singleton.process_video(
                input_video_path=refined_silent_video_path,
                output_video_path=hq_silent_video_path,
                prompt=global_prompt
            )
        except Exception as e:
            logger.error(f"Falha no processo de aprimoramento HD. Usando o vídeo refinado como fallback. Erro: {e}")
            # Se o HD falhar, usamos o vídeo refinado (silencioso) como base para o final
            os.rename(refined_silent_video_path, hq_silent_video_path)

        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Finalizando montagem...")
        final_video_path = os.path.join(self.workspace_dir, f"{base_name}_FINAL.mp4")

        #if audio_path and os.path.exists(audio_path):
        #    # Se o áudio foi gerado, combina o vídeo de ALTA QUALIDADE com ele
        #    self._combine_video_and_audio_ffmpeg(hq_silent_video_path, audio_path, final_video_path)
        #else:
        #    # Se não houver áudio, apenas renomeia o vídeo de alta qualidade
        #    os.rename(hq_silent_video_path, final_video_path)

        logger.info(f"Processo concluído! Vídeo final salvo em: {hq_silent_video_path}")
        yield {"final_path": hq_silent_video_path}

    def refine_latents(self, latents: torch.Tensor,
                       fps: int = 24,
                       denoise_strength: float = 0.35,
                       refine_steps: int = 12,
                       motion_prompt: str = "refining video, improving details, cinematic quality") -> torch.Tensor:
        """Aplica um passe de refinamento (denoise) em um tensor latente."""
        logger.info(f"Refinando tensor latente com shape {latents.shape}.")
        _, _, num_frames, latent_h, latent_w = latents.shape
        vae_scale_factor = self.vae.config.scaling_factor if hasattr(self.vae.config, 'scaling_factor') else 8
        pixel_height, pixel_width = latent_h * vae_scale_factor, latent_w * vae_scale_factor

        refined_latents_tensor, _ = self.ltx_manager.refine_latents(
            latents, height=pixel_height, width=pixel_width, video_total_frames=num_frames,
            video_fps=fps, motion_prompt=motion_prompt, current_fragment_index=int(time.time()),
            denoise_strength=denoise_strength, refine_steps=refine_steps)
        
        return refined_latents_tensor

    def upscale_latents(self, latents: torch.Tensor) -> torch.Tensor:
        """Interface para o UpscalerSpecialist."""
        logger.info(f"Realizando upscale em tensor latente com shape {latents.shape}.")
        return upscaler_specialist_singleton.upscale(latents)
        
    def _generate_latent_tensor_internal(self, conditioning_items, ltx_params, target_resolution, total_frames_to_generate):
        kwargs = {
            **ltx_params, 'width': target_resolution[1], 'height': target_resolution[0],
            'video_total_frames': total_frames_to_generate, 'video_fps': 24,
            'current_fragment_index': int(time.time()), 'conditioning_items_data': conditioning_items
        }
        return self.ltx_manager.generate_latent_fragment(**kwargs)

    def _quantize_to_multiple(self, n, m):
        if m == 0: return n
        quantized = int(round(n / m) * m)
        return m if n > 0 and quantized == 0 else quantized