Update deformes4D_engine.py

deformes4D_engine.py

@@ -102,7 +102,6 @@ class Deformes4DEngine:
             current_step += 1
             progress(current_step / TOTAL_STEPS, desc=f"Gerando Fragmento Bruto {fragment_index}/{num_transitions_to_generate}")
             
-            # ... (Lógica do Gemini para obter motion_prompt - 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]
@@ -119,7 +118,6 @@ class Deformes4DEngine:
             downscaled_width = self._quantize_to_multiple(int(expected_width * downscale_factor), 8)
             target_resolution_tuple = (downscaled_height, downscaled_width)
             
-            # ... (Lógica de condicionamento - sem alterações) ...
             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)
@@ -133,7 +131,6 @@ class Deformes4DEngine:
             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)
             
-            # ... (Lógica de poda e extração de Eco/Déjà-Vu - sem alterações) ...
             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()
@@ -148,13 +145,18 @@ class Deformes4DEngine:
         # --- ATO II: PÓS-PRODUÇÃO LATENTE GLOBAL ---
         current_step += 1
         progress(current_step / TOTAL_STEPS, desc="Unificação Causal (Concatenação)...")
-        concatenated_raw_latents = torch.cat(raw_latent_fragments, dim=2)
-        del raw_latent_fragments; gc.collect(); torch.cuda.empty_cache()
+        
+        # --- [INÍCIO DA CORREÇÃO MULTI-GPU] ---
+        logger.info(f"Movendo {len(raw_latent_fragments)} fragmentos latentes para o dispositivo principal ({self.device}) para unificação.")
+        tensors_on_main_device = [frag.to(self.device) for frag in raw_latent_fragments]
+        concatenated_raw_latents = torch.cat(tensors_on_main_device, dim=2)
+        # --- [FIM DA CORREÇÃO MULTI-GPU] ---
+
+        del raw_latent_fragments, tensors_on_main_device; gc.collect(); torch.cuda.empty_cache()
         logger.info(f"Latentes brutos unificados. Shape: {concatenated_raw_latents.shape}")
 
         current_step += 1
         progress(current_step / TOTAL_STEPS, desc="Polimento Global (Denoise)...")
-        # [ALTERAÇÃO CRÍTICA] Usamos um prompt vazio e guidance_scale=1.0 para um refinamento incondicional.
         denoised_latents = self.refine_latents(
             concatenated_raw_latents, 
             motion_prompt="", 
@@ -170,17 +172,15 @@ class Deformes4DEngine:
         refined_silent_video_path = os.path.join(self.workspace_dir, f"{base_name}_refined_silent.mp4")
 
         with imageio.get_writer(refined_silent_video_path, fps=FPS, codec='libx264', quality=8, output_params=['-pix_fmt', 'yuv420p']) as writer:
-            chunk_size = 7 # Tamanho de lote seguro para VRAM
+            chunk_size = 7
             latent_chunks = torch.split(denoised_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)}...")
                 
-                # Upscale e Decode por lote
                 upscaled_chunk = self.upscale_latents(latent_chunk)
                 pixel_tensor_chunk = self.latents_to_pixels(upscaled_chunk)
                 
-                # Converte e salva os frames
                 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)
@@ -207,7 +207,7 @@ class Deformes4DEngine:
                 prompt=global_prompt
             )
         except Exception as e:
-            logger.error(f"Falha no aprimoramento HD: {e}. Usando vídeo de qualidade padrão.")
+            logger.error(f"Falha no aprimoramento HD: {e}. Usando o vídeo de qualidade padrão.")
             os.rename(refined_silent_video_path, final_video_path)
 
         logger.info(f"Processo concluído! Vídeo final salvo em: {final_video_path}")
@@ -223,12 +223,11 @@ class Deformes4DEngine:
         pixel_width = latent_w * vae_scale_factor
         pixel_frames = (num_latent_frames - 1) * video_scale_factor
         
-        # [ALTERAÇÃO] Permite que guidance_scale seja passado como argumento
         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) # Usa o valor passado ou o padrão
+            "guidance_scale": kwargs.get('guidance_scale', 2.0)
         }
         
         refined_latents_tensor, _ = self.ltx_manager.refine_latents(latents, **final_ltx_params)