Update engineers/deformes3D.py

engineers/deformes3D.py

@@ -1,33 +1,17 @@
 # engineers/deformes3D.py
-# AducSdr: Uma implementação aberta e funcional da arquitetura ADUC-SDR
-# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
 #
-# Contato:
-# Carlos Rodrigues dos Santos
-# [email protected]
-# Rua Eduardo Carlos Pereira, 4125, B1 Ap32, Curitiba, PR, Brazil, CEP 8102025
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Repositórios e Projetos Relacionados:
-# GitHub: https://github.com/carlex22/Aduc-sdr
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+# Version: 1.5.0
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License...
 # PENDING PATENT NOTICE: Please see NOTICE.md.
 #
-# Version 1.4.5
+# This version implements an experimental dual-generation workflow. For each
+# keyframe, it first generates a version using the FLUX specialist, then
+# generates a second, "enriched" version using the LTX video engine to
+# allow for direct comparison of the models' visual languages.
 
 from PIL import Image
 import os
@@ -35,16 +19,21 @@ import time
 import logging
 import gradio as gr
 import yaml
+import torch
+import numpy as np
 
 from managers.flux_kontext_manager import flux_kontext_singleton
 from engineers.deformes2D_thinker import deformes2d_thinker_singleton
+from aduc_types import LatentConditioningItem
+from managers.ltx_manager import ltx_manager_singleton
+from managers.vae_manager import vae_manager_singleton
+from managers.latent_enhancer_manager import latent_enhancer_specialist_singleton
 
 logger = logging.getLogger(__name__)
 
 class Deformes3DEngine:
     """
     ADUC Specialist for static image (keyframe) generation.
-    This is responsible for the entire process of turning a script into a gallery of keyframes.
     """
     def __init__(self, workspace_dir):
         self.workspace_dir = workspace_dir
@@ -53,7 +42,7 @@ class Deformes3DEngine:
 
     def _generate_single_keyframe(self, prompt: str, reference_images: list[Image.Image], output_filename: str, width: int, height: int, callback: callable = None) -> str:
         """
-        Low-level function that generates a single image.
+        Low-level function that generates a single image using the FLUX helper.
         """
         logger.info(f"Generating keyframe '{output_filename}' with prompt: '{prompt}'")
         generated_image = self.image_generation_helper.generate_image(
@@ -67,55 +56,116 @@ class Deformes3DEngine:
 
     def generate_keyframes_from_storyboard(self, storyboard: list, initial_ref_path: str, global_prompt: str, keyframe_resolution: int, general_ref_paths: list, progress_callback_factory: callable = None):
         """
-        Orchestrates the generation of all keyframes from a storyboard.
+        Orchestrates the generation of all keyframes. For each keyframe, first
+        generates a version with FLUX, and then an "enriched" version with LTX
+        for direct comparison.
         """
         current_base_image_path = initial_ref_path
         previous_prompt = "N/A (initial reference image)"
-        final_keyframes = [current_base_image_path]
+        final_keyframes_gallery = [current_base_image_path]
         width, height = keyframe_resolution, keyframe_resolution
+        target_resolution_tuple = (width, height)
         
         num_keyframes_to_generate = len(storyboard) - 1
-        
-        logger.info(f"IMAGE SPECIALIST: Received order to generate {num_keyframes_to_generate} keyframes.")
+        logger.info(f"IMAGE SPECIALIST: Received order to generate {num_keyframes_to_generate} keyframes (FLUX + LTX versions).")
 
         for i in range(num_keyframes_to_generate):
+            scene_index = i + 1
             current_scene = storyboard[i]
             future_scene = storyboard[i+1]
-            progress_callback = progress_callback_factory(i + 1, num_keyframes_to_generate) if progress_callback_factory else None
+            progress_callback_flux = progress_callback_factory(scene_index, num_keyframes_to_generate) if progress_callback_factory else None
             
-            logger.info(f"--> Generating Keyframe {i+1}/{num_keyframes_to_generate}...")
+            logger.info(f"--> Generating Keyframe {scene_index}/{num_keyframes_to_generate}...")
+
+            # --- STEP A: Generate with FLUX ---
+            logger.info(f"    - Step A: Generating with FLUX...")
             
-            new_flux_prompt = deformes2d_thinker_singleton.get_anticipatory_keyframe_prompt(
+            flux_prompt = deformes2d_thinker_singleton.get_anticipatory_keyframe_prompt(
                 global_prompt=global_prompt, scene_history=previous_prompt,
                 current_scene_desc=current_scene, future_scene_desc=future_scene,
                 last_image_path=current_base_image_path, fixed_ref_paths=general_ref_paths
             )
             
-            images_for_flux_paths = list(set([current_base_image_path] + general_ref_paths))
-            images_for_flux = [Image.open(p) for p in images_for_flux_paths]
+            flux_ref_paths = list(set([current_base_image_path] + general_ref_paths))
+            flux_ref_images = [Image.open(p) for p in flux_ref_paths]
             
-            new_keyframe_path = self._generate_single_keyframe(
-                prompt=new_flux_prompt, reference_images=images_for_flux,
-                output_filename=f"keyframe_{i+1}.png", width=width, height=height,
-                callback=progress_callback
+            flux_keyframe_path = self._generate_single_keyframe(
+                prompt=flux_prompt, reference_images=flux_ref_images,
+                output_filename=f"keyframe_{scene_index}_flux.png", width=width, height=height,
+                callback=progress_callback_flux
             )
+            final_keyframes_gallery.append(flux_keyframe_path)
+            
+            # --- STEP B: LTX Enrichment Experiment ---
+            logger.info(f"    - Step B: Generating enrichment with LTX...")
 
-            final_keyframes.append(new_keyframe_path)
-            current_base_image_path = new_keyframe_path
-            previous_prompt = new_flux_prompt
+            ltx_conditioning_items = []
+            context_paths = [current_base_image_path] + [p for p in general_ref_paths if p != current_base_image_path][:3]
             
-        logger.info(f"IMAGE SPECIALIST: Keyframe generation complete.")
-        return final_keyframes
+            weight = 1.0
+            for idx, path in enumerate(context_paths):
+                img_pil = Image.open(path).convert("RGB")
+                img_processed = self._preprocess_image_for_latent_conversion(img_pil, target_resolution_tuple)
+                pixel_tensor = self._pil_to_pixel_tensor(img_processed)
+                latent_tensor = vae_manager_singleton.encode(pixel_tensor)
+                
+                ltx_conditioning_items.append(LatentConditioningItem(latent_tensor, 0, weight))
+                
+                if idx >= 0:
+                    weight -= 0.2
+            
+            ltx_base_params = {"guidance_scale": 3.0, "stg_scale": 0.1, "num_inference_steps": 25}
+            generated_latents, _ = ltx_manager_singleton.generate_latent_fragment(
+                height=height, width=width,
+                conditioning_items_data=ltx_conditioning_items,
+                motion_prompt=flux_prompt,
+                video_total_frames=16,
+                video_fps=24,
+                **ltx_base_params
+            )
+
+            final_latent = generated_latents[:, :, -1:, :, :]
+            upscaled_latent = latent_enhancer_specialist_singleton.upscale(final_latent)
+            enriched_pixel_tensor = vae_manager_singleton.decode(upscaled_latent)
+
+            ltx_keyframe_path = os.path.join(self.workspace_dir, f"keyframe_{scene_index}_ltx.png")
+            self.save_image_from_tensor(enriched_pixel_tensor, ltx_keyframe_path)
+            final_keyframes_gallery.append(ltx_keyframe_path)
+            
+            current_base_image_path = flux_keyframe_path 
+            previous_prompt = flux_prompt
+
+        logger.info(f"IMAGE SPECIALIST: Generation of all keyframe versions (FLUX + LTX) complete.")
+        return final_keyframes_gallery
+
+    # --- HELPER FUNCTIONS ---
+
+    def _preprocess_image_for_latent_conversion(self, image: Image.Image, target_resolution: tuple) -> Image.Image:
+        """Resizes and fits an image to the target resolution for VAE encoding."""
+        if image.size != target_resolution:
+            return ImageOps.fit(image, target_resolution, Image.Resampling.LANCZOS)
+        return image
+        
+    def _pil_to_pixel_tensor(self, pil_image: Image.Image) -> torch.Tensor:
+        """Helper to convert PIL to the 5D pixel tensor the VAE expects."""
+        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)
+        return (tensor * 2.0) - 1.0
+
+    def save_image_from_tensor(self, pixel_tensor: torch.Tensor, path: str):
+        """Helper to save a 1-frame pixel tensor as an image."""
+        tensor_chw = pixel_tensor.squeeze(0).squeeze(1)
+        tensor_hwc = tensor_chw.permute(1, 2, 0)
+        tensor_hwc = (tensor_hwc.clamp(-1, 1) + 1) / 2.0
+        image_np = (tensor_hwc.cpu().float().numpy() * 255).astype(np.uint8)
+        Image.fromarray(image_np).save(path)
 
 # --- Singleton Instantiation ---
 try:
     with open("config.yaml", 'r') as f:
         config = yaml.safe_load(f)
     WORKSPACE_DIR = config['application']['workspace_dir']
-    
-    # Correctly instantiate the Deformes3DEngine class
     deformes3d_engine_singleton = Deformes3DEngine(workspace_dir=WORKSPACE_DIR)
-    
 except Exception as e:
     logger.error(f"Could not initialize Deformes3DEngine: {e}", exc_info=True)
     deformes3d_engine_singleton = None