Upload pipeline_ltx_video.py

LTX-Video/ltx_video/pipelines/pipeline_ltx_video.py

@@ -45,17 +45,8 @@ from ltx_video.models.autoencoders.vae_encode import (
 )
 
 
-import warnings
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=FutureWarning)
-warnings.filterwarnings("ignore", message=".*")
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
-from huggingface_hub import logging as ll
-
-ll.set_verbosity_error()
-ll.set_verbosity_warning()
-ll.set_verbosity_info()
-ll.set_verbosity_debug()
 
 ASPECT_RATIO_1024_BIN = {
     "0.25": [512.0, 2048.0],
@@ -1389,42 +1380,6 @@ class LTXVideoPipeline(DiffusionPipeline):
         tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
         return torch.where(tokens_to_denoise_mask, denoised_latents, latents)
 
-    def _prepare_conditioning_media_to_latents(
-        self,
-        conditioning_item: ConditioningItem,
-        height: int,
-        width: int,
-        latent_height: int,
-        latent_width: int,
-        vae_per_channel_normalize: bool,
-        init_latents_dtype: torch.dtype,
-        init_latents_device: torch.device,
-    ) -> Tuple[ConditioningItem, torch.Tensor]:
-
-        media_item = conditioning_item.media_item
-        c = media_item.shape[1]
-
-        print (f"_ltx_prepare_media_item.shape {media_item.shape}")
-
-        # Dentro de _prepare_conditioning_media_to_latents:
-        c = media_item.shape[1]
-        if c == self.transformer.config.in_channels:
-            latents = media_item.to(dtype=init_latents_dtype, device=init_latents_device)
-            return conditioning_item, latents
-
-        conditioning_item = self._resize_conditioning_item(conditioning_item, height, width)
-        media_item = conditioning_item.media_item
-        latents = vae_encode(
-            media_item.to(dtype=self.vae.dtype, device=self.vae.device),
-            self.vae,
-            vae_per_channel_normalize=vae_per_channel_normalize,
-        ).to(dtype=init_latents_dtype, device=init_latents_device)
-        
-        print (f"_ltx_prepare_media_item_vae.shape? {media_item.shape}")
-        
-        return conditioning_item, latents
-
-
     def prepare_conditioning(
         self,
         conditioning_items: Optional[List[ConditioningItem]],
@@ -1434,119 +1389,6 @@ class LTXVideoPipeline(DiffusionPipeline):
         width: int,
         vae_per_channel_normalize: bool = False,
         generator=None,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
-        assert isinstance(self.vae, CausalVideoAutoencoder)
-        if not conditioning_items:
-            # comportamento existente
-            # ...
-            return init_latents, init_pixel_coords, None, 0
-    
-        # Deriva dims latentes
-        latent_height = height // self.vae_scale_factor
-        latent_width  = width  // self.vae_scale_factor
-    
-        batch_size, _, num_latent_frames = init_latents.shape[:3]
-        init_conditioning_mask = torch.zeros(
-            init_latents[:, 0, :, :, :].shape,
-            dtype=torch.float32,
-            device=init_latents.device,
-        )
-    
-        extra_conditioning_latents = []
-        extra_conditioning_pixel_coords = []
-        extra_conditioning_mask = []
-        extra_conditioning_num_latents = 0
-    
-        for conditioning_item in conditioning_items:
-            # NOVO: centraliza resize/encode/deteção de latents
-            conditioning_item, media_item_latents = self._prepare_conditioning_media_to_latents(
-                conditioning_item=conditioning_item,
-                height=height,
-                width=width,
-                latent_height=latent_height,
-                latent_width=latent_width,
-                vae_per_channel_normalize=vae_per_channel_normalize,
-                init_latents_dtype=init_latents.dtype,
-                init_latents_device=init_latents.device,
-            )
-    
-            media_frame_number = conditioning_item.media_frame_number
-            strength = conditioning_item.conditioning_strength
-    
-            # Validações de shape/frames (agora sobre latents)
-            b, c_l, f_l, h_l, w_l = media_item_latents.shape
-            assert c_l == self.transformer.config.in_channels
-            assert (h_l, w_l) == (latent_height, latent_width), "Latents com HxW incompatíveis"
-            assert f_l % 8 == 1, "n_frames latente deve satisfazer múltiplos da escala temporal + 1"
-            assert media_frame_number >= 0 and (media_frame_number + f_l) <= num_frames
-    
-            print(f"media_item_latents, {media_item_latents.shape}")
-    
-            # A partir daqui permanece igual (posicionamento, máscaras, patchify etc.)
-            if media_frame_number == 0:
-                media_item_latents, l_x, l_y = self._get_latent_spatial_position(
-                    media_item_latents, conditioning_item, height, width, strip_latent_border=True
-                )
-                init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = torch.lerp(
-                    init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l],
-                    media_item_latents,
-                    strength,
-                )
-                init_conditioning_mask[:, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = strength
-            else:
-                if f_l > 1:
-                    init_latents, init_conditioning_mask, media_item_latents = self._handle_non_first_conditioning_sequence(
-                        init_latents, init_conditioning_mask, media_item_latents, media_frame_number, strength,
-                    )
-                if media_item_latents is not None:
-                    noise = randn_tensor(
-                        media_item_latents.shape, generator=generator,
-                        device=media_item_latents.device, dtype=media_item_latents.dtype,
-                    )
-                    media_item_latents = torch.lerp(noise, media_item_latents, strength)
-                    media_item_latents, latent_coords = self.patchifier.patchify(latents=media_item_latents)
-                    pixel_coords = latent_to_pixel_coords(
-                        latent_coords, self.vae,
-                        causal_fix=self.transformer.config.causal_temporal_positioning,
-                    )
-                    pixel_coords[:, 0] += media_frame_number
-                    extra_conditioning_num_latents += media_item_latents.shape[1]
-                    conditioning_mask = torch.full(
-                        media_item_latents.shape[:2], strength,
-                        dtype=torch.float32, device=init_latents.device,
-                    )
-                    extra_conditioning_latents.append(media_item_latents)
-                    extra_conditioning_pixel_coords.append(pixel_coords)
-                    extra_conditioning_mask.append(conditioning_mask)
-    
-        # Patchify e concat iguais ao código existente...
-        init_latents, init_latent_coords = self.patchifier.patchify(latents=init_latents)
-        init_pixel_coords = latent_to_pixel_coords(
-            init_latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning,
-        )
-        init_conditioning_mask, _ = self.patchifier.patchify(latents=init_conditioning_mask.unsqueeze(1))
-        init_conditioning_mask = init_conditioning_mask.squeeze(-1)
-        if extra_conditioning_latents:
-            init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
-            init_pixel_coords = torch.cat([*extra_conditioning_pixel_coords, init_pixel_coords], dim=2)
-            init_conditioning_mask = torch.cat([*extra_conditioning_mask, init_conditioning_mask], dim=1)
-        if self.transformer.use_tpu_flash_attention:
-            init_latents = init_latents[:, :-extra_conditioning_num_latents]
-            init_pixel_coords = init_pixel_coords[:, :, :-extra_conditioning_num_latents]
-            init_conditioning_mask = init_conditioning_mask[:, :-extra_conditioning_num_latents]
-        return init_latents, init_pixel_coords, init_conditioning_mask, extra_conditioning_num_latents
-        
-
-    
-    def prepare_conditioning_old(
-        self,
-        conditioning_items: Optional[List[ConditioningItem]],
-        init_latents: torch.Tensor,
-        num_frames: int,
-        height: int,
-        width: int,
-        vae_per_channel_normalize: bool = False,
-        generator=None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
         """
         Prepare conditioning tokens based on the provided conditioning items.
@@ -1596,33 +1438,41 @@ class LTXVideoPipeline(DiffusionPipeline):
 
             # Process each conditioning item
             for conditioning_item in conditioning_items:
-                conditioning_item = self._resize_conditioning_item(
-                    conditioning_item, height, width
-                )
-                media_item = conditioning_item.media_item
-                media_frame_number = conditioning_item.media_frame_number
-                strength = conditioning_item.conditioning_strength
-                assert media_item.ndim == 5  # (b, c, f, h, w)
-                b, c, n_frames, h, w = media_item.shape
-                assert (
-                    height == h and width == w
-                ) or media_frame_number == 0, f"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0"
-                assert n_frames % 8 == 1
-                assert (
-                    media_frame_number >= 0
-                    and media_frame_number + n_frames <= num_frames
-                )
-
-                # Encode the provided conditioning media item
-                media_item_latents = vae_encode(
-                    media_item.to(dtype=self.vae.dtype, device=self.vae.device),
-                    self.vae,
-                    vae_per_channel_normalize=vae_per_channel_normalize,
-                ).to(dtype=init_latents.dtype)
-
-                print(f"media_item_latents, {media_item_latents.shape}")
-
+            
+                print(f"media_item_latents ini {conditioning_item.media_item.shape}")
                 
+                c = conditioning_item.media_item.shape[1]
+                if c == self.transformer.config.in_channels:
+                    media_item_latents = conditioning_item.media_item.to(dtype=init_latents_dtype, device=init_latents_device)
+                    strength = conditioning_item.conditioning_strength
+                    media_frame_number = conditioning_item.media_frame_number
+                else: 
+                    conditioning_item = self._resize_conditioning_item(
+                        conditioning_item, height, width
+                    )
+                    media_item = conditioning_item.media_item
+                    media_frame_number = conditioning_item.media_frame_number
+                    strength = conditioning_item.conditioning_strength
+                    assert media_item.ndim == 5  # (b, c, f, h, w)
+                    b, c, n_frames, h, w = media_item.shape
+                    assert (
+                        height == h and width == w
+                    ) or media_frame_number == 0, f"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0"
+                    assert n_frames % 8 == 1
+                    assert (
+                        media_frame_number >= 0
+                        and media_frame_number + n_frames <= num_frames
+                    )
+    
+                    # Encode the provided conditioning media item
+                    media_item_latents = vae_encode(
+                        media_item.to(dtype=self.vae.dtype, device=self.vae.device),
+                        self.vae,
+                        vae_per_channel_normalize=vae_per_channel_normalize,
+                    ).to(dtype=init_latents.dtype)
+                    
+                    print(f"media_item_latents encode vae {conditioning_item.media_item.shape}")
+   
                 # Handle the different conditioning cases
                 if media_frame_number == 0:
                     # Get the target spatial position of the latent conditioning item
@@ -1763,7 +1613,6 @@ class LTXVideoPipeline(DiffusionPipeline):
         )
         return new_conditioning_item
 
-    @staticmethod
     def _get_latent_spatial_position(
         self,
         latents: torch.Tensor,
@@ -1887,7 +1736,6 @@ class LTXVideoPipeline(DiffusionPipeline):
             latents,
         )
 
-    @staticmethod
     def trim_conditioning_sequence(
         self, start_frame: int, sequence_num_frames: int, target_num_frames: int
     ):