Update aduc_framework/managers/wan_manager.py

aduc_framework/managers/wan_manager.py

@@ -1,138 +1,174 @@
-# aduc_framework/managers/vae_wan_manager.py (Versão Definitiva Completa)
+# aduc_framework/managers/wan_manager.py (Versão Definitiva Completa)
 import torch
 import logging
 import yaml
 from PIL import Image
-import numpy as np
 from typing import List, Optional
 import sys
 import os
 
 # --- INÍCIO DA CORREÇÃO DE IMPORTAÇÃO ---
-# Adiciona o diretório do Wan2.2 ao sys.path para encontrar módulos customizados.
+# Adiciona o diretório do Wan2.2 ao sys.path para que o Python o encontre.
 WAN_REPO_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'Wan2.2'))
 if WAN_REPO_PATH not in sys.path:
     sys.path.insert(0, WAN_REPO_PATH)
-    logging.info(f"Adicionado '{WAN_REPO_PATH}' ao sys.path para o VaeWanManager.")
+    logging.info(f"Adicionado '{WAN_REPO_PATH}' ao sys.path para importações do WanManager.")
 # --- FIM DA CORREÇÃO DE IMPORTAÇÃO ---
 
-# Ferramentas da nossa arquitetura
+# Ferramentas da nossa arquitetura ADUC
 from ..tools.hardware_manager import hardware_manager
+from ..tools.pipeline_patches import apply_aduc_patches
+from ..types import LatentConditioningItem
 
-# --- IMPORTAÇÕES CORRIGIDAS ---
-# Importa a classe de VAE customizada do Wan2.2 e a função auxiliar oficial.
-from diffusers.models import AutoencoderKLWan
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
+# Especialistas e modelos necessários
+from .vae_wan_manager import vae_wan_manager_singleton
+#from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+#from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
+from transformers import CLIPVisionModel
 
-logger = logging.getLogger(__name__)
 
+from diffusers import WanImageToVideoPipeline
+#from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.models import WanTransformer3DModel
+
+logger = logging.getLogger(__name__)
 
-class VaeWanManager:
+class WanManager:
     """
-    Especialista VAE dedicado e persistente para a pipeline Wan2.2.
-    Gerencia o ciclo de vida do AutoencoderKLWan em uma GPU dedicada,
-    garantindo a tradução correta entre o espaço de pixels e o espaço latente.
+    Especialista ADUC completo e otimizado para geração de vídeo com Wan2.2.
+    Incorpora a fusão do LoRA Lightning para geração de alta velocidade (8-steps)
+    e patches customizados para controle temporal preciso.
     """
     def __init__(self):
         self.device = None
-        self.vae: Optional[AutoencoderKLWan] = None
-        self.dtype = None
+        self.pipe: Optional[WanImageToVideoPipeline] = None
         self.config = self._load_config()
-
+        
         if self.config:
             gpus_required = self.config.get('gpus_required', 0)
             if gpus_required > 0:
-                self.device = hardware_manager.allocate_gpus('VaeWanManager', gpus_required)[0]
-                logger.info(f"VaeWanManager: GPU dedicada '{self.device}' reservada.")
+                self.device = hardware_manager.allocate_gpus('WanManager', gpus_required)[0]
+                logger.info(f"WanManager (Lightning): GPU {self.device} reservada.")
             else:
-                self.device = torch.device('cpu')
-                logger.warning("VaeWanManager: Nenhuma GPU dedicada foi alocada.")
+                logger.warning("WanManager está desabilitado (gpus_required: 0).")
         else:
-            logger.warning("Configuração para 'vae_wan' não encontrada em config.yaml.")
+            logger.warning("Configuração para 'wan' não encontrada.")
 
     def _load_config(self):
-        """Carrega a configuração específica deste manager."""
+        """Carrega a configuração específica deste manager do arquivo YAML global."""
         try:
-            with open("config.yaml", 'r', encoding='utf-8') as f:
-                return yaml.safe_load(f).get('specialists', {}).get('vae_wan', {})
+            with open("config.yaml", 'r') as f:
+                return yaml.safe_load(f).get('specialists', {}).get('wan', {})
         except FileNotFoundError:
             logger.error("config.yaml não encontrado.")
             return None
 
     def _lazy_init(self):
-        """Carrega o modelo VAE do Wan2.2 para a VRAM no primeiro uso."""
-        if self.vae is not None:
+        """Carrega a pipeline, aplica otimizações, funde o LoRA e aplica nosso patch."""
+        if self.pipe is not None:
             return
         if not self.device or not self.config:
-            raise RuntimeError("VaeWanManager não pode ser inicializado.")
+            raise RuntimeError("WanManager não pode ser inicializado.")
 
-        logger.info(f"VAE-WAN MANAGER ({self.device}): Carregando VAE do Wan2.2...")
+        logger.info(f"WAN MANAGER ({self.device}): Iniciando carregamento OTIMIZADO do Wan2.2...")
         
-        model_id = self.config.get("model_id")
-        local_model_path = f"/app/models/{model_id}" # Carrega do nosso diretório local
-        self.dtype = torch.float32 # VAEs são mais estáveis em FP32
+        main_model_id = self.config.get("model_id")
+        opt_model_id = self.config.get("optimized_model_id")
+        lora_repo = self.config.get("lora_repo")
+        lora_filename = self.config.get("lora_filename")
+        torch_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+        
+        # Define os caminhos locais onde os modelos foram baixados
+        main_model_path = f"/app/models/{main_model_id}"
+        opt_model_path = f"/app/models/{opt_model_id}"
+        lora_path = f"/app/models/loras/{os.path.basename(lora_filename)}"
 
         try:
-            # Carrega a classe CORRETA (AutoencoderKLWan) do caminho LOCAL.
-            self.vae = AutoencoderKLWan.from_pretrained(
-                local_model_path, subfolder="vae", torch_dtype=self.dtype
-            ).to(self.device)
-            self.vae.eval()
-            logger.info(f"VAE-WAN MANAGER ({self.device}): VAE do Wan2.2 pronto e 'quente' na VRAM.")
+            # 1. Requisição do VAE dedicado
+            vae_wan_manager_singleton._lazy_init()
+            vae = vae_wan_manager_singleton.vae
+            if vae is None: raise RuntimeError("Falha ao obter o VAE do vae_wan_manager_singleton.")
+
+            # 2. Carregamento dos componentes dos caminhos locais
+            image_encoder = CLIPVisionModel.from_pretrained(main_model_path, subfolder="image_encoder")
+            transformer = WanTransformer3DModel.from_pretrained(opt_model_path, subfolder='transformer', torch_dtype=torch_dtype)
+            transformer_2 = WanTransformer3DModel.from_pretrained(opt_model_path, subfolder='transformer_2', torch_dtype=torch_dtype)
+
+            # 3. Montagem da pipeline base
+            self.pipe = WanImageToVideoPipeline.from_pretrained(main_model_path, vae=vae, image_encoder=image_encoder, transformer=transformer, transformer_2=transformer_2, torch_dtype=torch_dtype)
+            
+            # 4. Ajuste do Scheduler
+            self.pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(self.pipe.scheduler.config, shift=32.0)
+            
+            # 5. Fusão do LoRA Lightning
+            logger.info(f"WAN MANAGER ({self.device}): Carregando e fundindo LoRA Lightning de '{lora_path}'...")
+            self.pipe.load_lora_weights(os.path.dirname(lora_path), weight_name=os.path.basename(lora_path), adapter_name="lightx2v")
+            self.pipe.load_lora_weights(os.path.dirname(lora_path), weight_name=os.path.basename(lora_path), adapter_name="lightx2v_2", load_into_transformer_2=True)
+            self.pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1.0, 1.0])
+            self.pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3.0, components=["transformer"])
+            self.pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1.0, components=["transformer_2"])
+            self.pipe.unload_lora_weights()
+            logger.info(f"WAN MANAGER ({self.device}): LoRA Lightning fundido.")
+
+            # 6. Aplicação do nosso patch ADUC
+            apply_aduc_patches()
+            
+            # 7. Finalização e envio para a GPU
+            self.pipe.to(self.device)
+            logger.info(f"WAN MANAGER ({self.device}): Pipeline Wan2.2 OTIMIZADA, MODIFICADA e pronta na VRAM.")
+
         except Exception as e:
-            logger.error(f"VAE-WAN MANAGER: Falha CRÍTICA ao carregar o VAE: {e}", exc_info=True)
-            self.vae = None
+            logger.error(f"WAN MANAGER: Falha CRÍTICA ao carregar a pipeline: {e}", exc_info=True)
+            self.pipe = None
             raise e
 
-    def _preprocess_pil_image(self, pil_image: Image.Image, target_resolution: tuple) -> torch.Tensor:
-        """Converte uma imagem PIL para o formato de tensor 5D esperado pelo VAE de vídeo."""
-        from PIL import ImageOps
-        img = pil_image.convert("RGB")
-        processed_img = ImageOps.fit(img, target_resolution, Image.Resampling.LANCZOS)
-        image_np = np.array(processed_img).astype(np.float32) / 255.0
-        
-        # Converte para (B, C, H, W)
-        tensor_4d = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0)
-        tensor_4d_normalized = (tensor_4d * 2.0) - 1.0
-        
-        # Adiciona a dimensão de "frame" para criar um tensor 5D (B, C, F, H, W)
-        tensor_5d = tensor_4d_normalized.unsqueeze(2)
-        
-        return tensor_5d
-
     @torch.no_grad()
-    def encode_batch(self, pil_images: List[Image.Image], target_resolution: tuple) -> List[torch.Tensor]:
-        """Codifica um lote de imagens PIL para o espaço latente do Wan."""
+    def generate_latent_fragment(self, **kwargs) -> tuple[torch.Tensor, None]:
+        """Gera um fragmento de vídeo no espaço latente. A interface é IDÊNTICA à do LtxManager."""
         self._lazy_init()
-        if not self.vae:
-            raise RuntimeError("O VAE do WanManager não foi carregado.")
+        if not self.pipe:
+            raise RuntimeError("A pipeline do WanManager não está disponível.")
+
+        conditioning_items: List[LatentConditioningItem] = kwargs.get("conditioning_items_data", [])
+        if not conditioning_items:
+            raise ValueError("WanManager no modo ADUC requer 'conditioning_items_data'.")
+
+        pipeline_params = {
+            "prompt": kwargs.get("motion_prompt", ""),
+            "negative_prompt": kwargs.get("negative_prompt", "static, disfigured, low quality"),
+            "height": kwargs.get("height", self.config.get("default_height", 480)),
+            "width": kwargs.get("width", self.config.get("default_width", 832)),
+            "num_frames": kwargs.get("video_total_frames", self.config.get("default_frames", 81)),
+            "guidance_scale": kwargs.get("guidance_scale", self.config.get("guidance_scale", 1.0)),
+            "guidance_scale_2": kwargs.get("guidance_scale_2", self.config.get("guidance_scale_2", 1.0)),
+            "num_inference_steps": kwargs.get("num_inference_steps", self.config.get("inference_steps", 8)),
+            "generator": torch.Generator(device=self.device).manual_seed(int(torch.randint(0, 100000, (1,)).item())),
+        }
         
-        latents_list = []
-        for img in pil_images:
-            # A função de pré-processamento agora retorna o tensor 5D correto
-            pixel_tensor_gpu = self._preprocess_pil_image(img, target_resolution).to(self.device, dtype=self.dtype)
-            
-            encoder_output = self.vae.encode(pixel_tensor_gpu)
-            latents = retrieve_latents(encoder_output)
-            
-            latents_list.append(latents.cpu())
-            
-        return latents_list
+        logger.info(f"WAN MANAGER (Lightning): Gerando fragmento com {pipeline_params['num_inference_steps']} passos.")
 
-    @torch.no_grad()
-    def decode(self, latent_tensor: torch.Tensor) -> torch.Tensor:
-        """Decodifica um tensor latente do Wan para o espaço de pixels."""
-        self._lazy_init()
-        if not self.vae:
-            raise RuntimeError("O VAE do WanManager não foi carregado.")
+        first_latent_for_encoding = conditioning_items[0].latent_tensor
+        pil_image_for_encoder = self._decode_latent_to_pil(first_latent_for_encoding)
 
-        latent_tensor_gpu = latent_tensor.to(self.device, dtype=self.dtype)
+        output = self.pipe(
+            image=pil_image_for_encoder,
+            conditioning_items=conditioning_items,
+            output_type='latent', 
+            **pipeline_params
+        )
         
-        # Acessa a saída através do atributo .sample
-        decode_output = self.vae.decode(latent_tensor_gpu)
-        pixels = decode_output.sample
+        video_latents = output.frames
+        logger.info(f"WAN MANAGER (Lightning): Fragmento latente gerado. Shape: {video_latents.shape}")
         
-        return pixels.cpu()
+        return video_latents.cpu(), None
+
+    def _decode_latent_to_pil(self, latent_tensor: torch.Tensor) -> Image.Image:
+        """Função auxiliar para decodificar um latente em uma imagem PIL usando o VAE dedicado."""
+        pixel_tensor = vae_wan_manager_singleton.decode(latent_tensor.unsqueeze(0))
+        pixel_tensor = (pixel_tensor / 2 + 0.5).clamp(0, 1)
+        numpy_image = (pixel_tensor.cpu().permute(0, 2, 3, 4, 1).squeeze(0).squeeze(0) * 255).byte().numpy()
+        return Image.fromarray(numpy_image)
 
 # --- Instância Singleton ---
-vae_wan_manager_singleton = VaeWanManager()
+wan_manager_singleton = WanManager()