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 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
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) 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 ---
    @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 _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)
    
    # --- 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, 
                            use_upscaler: bool, use_refiner: bool, use_hd: bool, use_audio: bool,
                            video_resolution: int, use_continuity_director: bool, 
                            progress: gr.Progress = gr.Progress()):
        
        num_transitions_to_generate = len(keyframes) - 1
        TOTAL_STEPS = num_transitions_to_generate + 4
        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)
        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": 2.0, "stg_scale": 0.025, "rescaling_scale": 0.15, "num_inference_steps": 20}
        
        # --- [INÍCIO DA CORREÇÃO] ---
        # Lógica robusta para extrair caminhos de arquivo da lista de keyframes.
        keyframe_paths = []
        for item in keyframes:
            if isinstance(item, str):
                keyframe_paths.append(item)
            elif isinstance(item, tuple) and len(item) > 0:
                keyframe_paths.append(item[0]) # Assume que o caminho está no primeiro elemento da tupla
            elif hasattr(item, 'name'):
                keyframe_paths.append(item.name)
            else:
                logger.warning(f"Item na lista de keyframes com tipo inesperado e sem atributo '.name': {type(item)}")
        # --- [FIM DA CORREÇÃO] ---

        story_history = ""
        eco_latent_for_next_loop = None
        dejavu_latent_for_next_loop = None
        raw_latent_fragments = []

        # --- ATO I: GERAÇÃO CAUSAL PURA (LOOP DE FRAGMENTOS) ---
        for i in range(num_transitions_to_generate):
            fragment_index = i + 1
            current_step += 1
            progress(current_step / TOTAL_STEPS, desc=f"Gerando Fragmento Causal {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 = 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 = video_resolution, video_resolution
            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

            raw_latent_fragments.append(latents_video)
        
        # --- ATO II: PÓS-PRODUÇÃO LATENTE GLOBAL (CONDICIONAL) ---
        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Unificação Causal (Concatenação)...")
        tensors_on_main_device = [frag.to(self.device) for frag in raw_latent_fragments]
        processed_latents = torch.cat(tensors_on_main_device, dim=2)
        del raw_latent_fragments, tensors_on_main_device; gc.collect(); torch.cuda.empty_cache()

        if use_refiner:
            current_step += 1
            progress(current_step / TOTAL_STEPS, desc="Polimento Global (Denoise)...")
            processed_latents = self.refine_latents(
                processed_latents, 
                motion_prompt="", 
                guidance_scale=1.0
            )
            logger.info(f"Polimento global aplicado. Shape: {processed_latents.shape}")
        else:
            logger.info("Etapa de refinamento desativada.")

        # --- ATO III: RENDERIZAÇÃO E FINALIZAÇÃO ---
        base_name = f"movie_{int(time.time())}"
        current_step += 1
        progress(current_step / TOTAL_STEPS, desc="Renderização (em lotes)...")
        intermediate_video_path = os.path.join(self.workspace_dir, f"{base_name}_intermediate.mp4")

        with imageio.get_writer(intermediate_video_path, fps=FPS, codec='libx264', quality=8, output_params=['-pix_fmt', 'yuv420p']) as writer:
            chunk_size = 15 if use_upscaler else 30
            latent_chunks = torch.split(processed_latents, chunk_size, dim=2)
            
            for i, latent_chunk in enumerate(latent_chunks):
                logger.info(f"Processando e renderizando lote {i+1}/{len(latent_chunks)}...")
                
                processed_chunk = self.upscale_latents(latent_chunk) if use_upscaler else latent_chunk
                pixel_tensor_chunk = self.latents_to_pixels(processed_chunk)
                
                pixel_tensor_chunk = pixel_tensor_chunk.squeeze(0).permute(1, 2, 3, 0)
                pixel_tensor_chunk = (pixel_tensor_chunk.clamp(-1, 1) + 1) / 2.0
                video_np_chunk = (pixel_tensor_chunk.detach().cpu().float().numpy() * 255).astype(np.uint8)
                
                for frame in video_np_chunk:
                    writer.append_data(frame)
                
                del latent_chunk, processed_chunk, pixel_tensor_chunk, video_np_chunk
                gc.collect()
                torch.cuda.empty_cache()

        del processed_latents; gc.collect(); torch.cuda.empty_cache()
        logger.info(f"Vídeo intermediário renderizado em: {intermediate_video_path}")

        final_video_path = os.path.join(self.workspace_dir, f"{base_name}_FINAL.mp4")
        
        if use_hd:
            current_step += 1
            progress(current_step / TOTAL_STEPS, desc="Masterização Final (HD)...")
            try:
                hd_specialist_singleton.process_video(
                    input_video_path=intermediate_video_path,
                    output_video_path=final_video_path,
                    prompt=global_prompt
                )
            except Exception as e:
                logger.error(f"Falha na masterização HD: {e}. Usando vídeo de qualidade padrão.")
                os.rename(intermediate_video_path, final_video_path)
        else:
            logger.info("Etapa de masterização HD desativada.")
            os.rename(intermediate_video_path, final_video_path)

        if use_audio:
            logger.warning("Geração de áudio solicitada, mas está desativada nesta versão do código.")

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

    def refine_latents(self, latents: torch.Tensor, fps: int = 24, denoise_strength: float = 0.35, refine_steps: int = 12, motion_prompt: str = "...", **kwargs) -> torch.Tensor:
        logger.info(f"Refinando tensor latente com shape {latents.shape}.")
        _, _, num_latent_frames, latent_h, latent_w = latents.shape
        video_scale_factor = getattr(self.vae.config, 'temporal_scale_factor', 8)
        vae_scale_factor = getattr(self.vae.config, 'spatial_downscale_factor', 8)
        
        pixel_height = latent_h * vae_scale_factor
        pixel_width = latent_w * vae_scale_factor
        pixel_frames = (num_latent_frames - 1) * video_scale_factor
        
        final_ltx_params = {
            "height": pixel_height, "width": pixel_width, "video_total_frames": pixel_frames,
            "video_fps": fps, "motion_prompt": motion_prompt, "current_fragment_index": int(time.time()),
            "denoise_strength": denoise_strength, "refine_steps": refine_steps,
            "guidance_scale": kwargs.get('guidance_scale', 2.0)
        }
        
        refined_latents_tensor, _ = self.ltx_manager.refine_latents(latents, **final_ltx_params)
        
        logger.info(f"Retornando tensor latente refinado com shape: {refined_latents_tensor.shape}")
        return refined_latents_tensor
        
    def upscale_latents(self, latents: torch.Tensor) -> torch.Tensor:
        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