Update engineers/deformes3D.py

engineers/deformes3D.py

@@ -2,28 +2,34 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 1.1.1
+# Version: 1.5.1
 #
-# This file defines the Deformes3DEngine, the specialist responsible for
-# generating the key visual anchors (keyframes) of the story. It acts as the
-# "Art Director," translating narrative scenes into static images.
+# This version maintains the core FLUX-based keyframe generation and adds the
+# LTX-based "enrichment" as a secondary, experimental step for each keyframe,
+# allowing for direct comparison without altering the primary workflow.
 
-from PIL import Image
+from PIL import Image, ImageOps
 import os
 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.
+    Primary engine is FLUX, with an experimental LTX enrichment step.
     """
     def __init__(self, workspace_dir):
         self.workspace_dir = workspace_dir
@@ -46,58 +52,121 @@ 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(scene_index, 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 {scene_index}/{num_keyframes_to_generate}...")
+
+            # --- STEP A: Generate with FLUX (Primary Method) ---
+            logger.info(f"    - Step A: Generating with FLUX...")
             
-            # Delegate the "thinking" part to the Deformes2DThinker
-            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]
             
-            # Execute the image generation
-            new_keyframe_path = self._generate_single_keyframe(
-                prompt=new_flux_prompt, reference_images=images_for_flux,
-                output_filename=f"keyframe_{scene_index}.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...")
+
+            ltx_context_paths = list(reversed(context_paths))
+            logger.info(f"    - LTX Context Order (Reversed): {[os.path.basename(p) for p in ltx_context_paths]}")
 
-            final_keyframes.append(new_keyframe_path)
-            current_base_image_path = new_keyframe_path
-            previous_prompt = new_flux_prompt
             
-        logger.info("IMAGE SPECIALIST: Keyframe generation complete.")
-        return final_keyframes
+            ltx_conditioning_items = []
+            context_paths = [current_base_image_path] + [p for p in general_ref_paths if p != current_base_image_path][:3]
+            
+            weight = 0.6
+            for idx, path in enumerate(ltx_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.1
+            
+            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)
+            
+            # Use the FLUX keyframe as the base for the next iteration to maintain the primary narrative 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