add docs tests + more refactor

Author: yiyixuxu

docs/source/en/api/pipelines/ltx_video.md

@@ -196,6 +196,12 @@ export_to_video(video, "ship.mp4", fps=24)
   - all
   - __call__
 
+## LTXConditionPipeline
+
+[[autodoc]] LTXConditionPipeline
+  - all
+  - __call__
+
 ## LTXPipelineOutput
 
 [[autodoc]] pipelines.ltx.pipeline_output.LTXPipelineOutput

scripts/convert_ltx_to_diffusers.py

@@ -105,6 +105,7 @@ def remove_keys_(key: str, state_dict: Dict[str, Any]):
     "per_channel_statistics.mean-of-means": remove_keys_,
     "per_channel_statistics.mean-of-stds": remove_keys_,
     "model.diffusion_model": remove_keys_,
+    "decoder.timestep_scale_multiplier": remove_keys_,
 }
 
 
@@ -270,6 +271,7 @@ def get_vae_config(version: str) -> Dict[str, Any]:
             "decoder_causal": False,
             "spatial_compression_ratio": 32,
             "temporal_compression_ratio": 8,
+            "timestep_scale_multiplier": 1000.0,
         }
         VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT)
     return config

src/diffusers/__init__.py

@@ -347,6 +347,7 @@
             "LDMTextToImagePipeline",
             "LEditsPPPipelineStableDiffusion",
             "LEditsPPPipelineStableDiffusionXL",
+            "LTXConditionPipeline",
             "LTXImageToVideoPipeline",
             "LTXPipeline",
             "Lumina2Text2ImgPipeline",
@@ -857,6 +858,7 @@
             LDMTextToImagePipeline,
             LEditsPPPipelineStableDiffusion,
             LEditsPPPipelineStableDiffusionXL,
+            LTXConditionPipeline,
             LTXImageToVideoPipeline,
             LTXPipeline,
             Lumina2Text2ImgPipeline,

src/diffusers/models/autoencoders/autoencoder_kl_ltx.py

@@ -921,12 +921,14 @@ def __init__(
         timestep_conditioning: bool = False,
         upsample_residual: Tuple[bool, ...] = (False, False, False, False),
         upsample_factor: Tuple[bool, ...] = (1, 1, 1, 1),
+        timestep_scale_multiplier: float = 1.0,
     ) -> None:
         super().__init__()
 
         self.patch_size = patch_size
         self.patch_size_t = patch_size_t
         self.out_channels = out_channels * patch_size**2
+        self.timestep_scale_multiplier = timestep_scale_multiplier
 
         block_out_channels = tuple(reversed(block_out_channels))
         spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
@@ -981,9 +983,7 @@ def __init__(
         # timestep embedding
         self.time_embedder = None
         self.scale_shift_table = None
-        self.timestep_scale_multiplier = None
         if timestep_conditioning:
-            self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0, dtype=torch.float32))
             self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(output_channel * 2, 0)
             self.scale_shift_table = nn.Parameter(torch.randn(2, output_channel) / output_channel**0.5)
 
@@ -992,7 +992,7 @@ def __init__(
     def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
         hidden_states = self.conv_in(hidden_states)
 
-        if self.timestep_scale_multiplier is not None:
+        if temb is not None:
             temb = temb * self.timestep_scale_multiplier
 
         if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -1107,6 +1107,7 @@ def __init__(
         decoder_causal: bool = False,
         spatial_compression_ratio: int = None,
         temporal_compression_ratio: int = None,
+        timestep_scale_multiplier: float = 1.0,
     ) -> None:
         super().__init__()
 
@@ -1137,6 +1138,7 @@ def __init__(
             inject_noise=decoder_inject_noise,
             upsample_residual=upsample_residual,
             upsample_factor=upsample_factor,
+            timestep_scale_multiplier=timestep_scale_multiplier,
         )
 
         latents_mean = torch.zeros((latent_channels,), requires_grad=False)

src/diffusers/pipelines/__init__.py

@@ -260,7 +260,7 @@
         ]
     )
     _import_structure["latte"] = ["LattePipeline"]
-    _import_structure["ltx"] = ["LTXPipeline", "LTXImageToVideoPipeline"]
+    _import_structure["ltx"] = ["LTXPipeline", "LTXImageToVideoPipeline", "LTXConditionPipeline"]
     _import_structure["lumina"] = ["LuminaText2ImgPipeline"]
     _import_structure["lumina2"] = ["Lumina2Text2ImgPipeline"]
     _import_structure["marigold"].extend(
@@ -610,7 +610,7 @@
             LEditsPPPipelineStableDiffusion,
             LEditsPPPipelineStableDiffusionXL,
         )
-        from .ltx import LTXImageToVideoPipeline, LTXPipeline
+        from .ltx import LTXConditionPipeline, LTXImageToVideoPipeline, LTXPipeline
         from .lumina import LuminaText2ImgPipeline
         from .lumina2 import Lumina2Text2ImgPipeline
         from .marigold import (

src/diffusers/pipelines/ltx/__init__.py

@@ -23,6 +23,7 @@
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
     _import_structure["pipeline_ltx"] = ["LTXPipeline"]
+    _import_structure["pipeline_ltx_condition"] = ["LTXConditionPipeline"]
     _import_structure["pipeline_ltx_image2video"] = ["LTXImageToVideoPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -34,6 +35,7 @@
         from ...utils.dummy_torch_and_transformers_objects import *
     else:
         from .pipeline_ltx import LTXPipeline
+        from .pipeline_ltx_condition import LTXConditionPipeline
         from .pipeline_ltx_image2video import LTXImageToVideoPipeline
 
 else:

src/diffusers/pipelines/ltx/pipeline_ltx_condition.py

@@ -21,6 +21,7 @@
 from transformers import T5EncoderModel, T5TokenizerFast
 
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...image_processor import PipelineImageInput
 from ...loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin
 from ...models.autoencoders import AutoencoderKLLTXVideo
 from ...models.transformers import LTXVideoTransformer3DModel
@@ -45,12 +46,11 @@
     Examples:
         ```py
         >>> import torch
-        >>> from diffusers import LTXImageToVideoPipeline
+        >>> from diffusers import LTXConditionPipeline
         >>> from diffusers.utils import export_to_video, load_image
 
-        >>> pipe = LTXImageToVideoPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
+        >>> pipe = LTXConditionPipeline.from_pretrained("YiYiXu/ltx-95", torch_dtype=torch.bfloat16)
         >>> pipe.to("cuda")
-
         >>> image = load_image(
         ...     "https://huggingface.co/datasets/a-r-r-o-w/tiny-meme-dataset-captioned/resolve/main/images/8.png"
         ... )
@@ -405,6 +405,11 @@ def encode_prompt(
     def check_inputs(
         self,
         prompt,
+        conditions,
+        image,
+        video,
+        frame_index,
+        strength,
         height,
         width,
         callback_on_step_end_tensor_inputs=None,
@@ -455,6 +460,26 @@ def check_inputs(
                     f" {negative_prompt_attention_mask.shape}."
                 )
 
+        if conditions is not None and (image is not None or video is not None):
+            raise ValueError("If `conditions` is provided, `image` and `video` must not be provided.")
+
+        if conditions is None and (image is None and video is None):
+            raise ValueError("If `conditions` is not provided, `image` or `video` must be provided.")
+
+        if conditions is None:
+            if isinstance(image, list) and isinstance(frame_index, list) and len(image) != len(frame_index):
+                raise ValueError(
+                    "If `conditions` is not provided, `image` and `frame_index` must be of the same length."
+                )
+            elif isinstance(image, list) and isinstance(strength, list) and len(image) != len(strength):
+                raise ValueError("If `conditions` is not provided, `image` and `strength` must be of the same length.")
+            elif isinstance(video, list) and isinstance(frame_index, list) and len(video) != len(frame_index):
+                raise ValueError(
+                    "If `conditions` is not provided, `video` and `frame_index` must be of the same length."
+                )
+            elif isinstance(video, list) and isinstance(strength, list) and len(video) != len(strength):
+                raise ValueError("If `conditions` is not provided, `video` and `strength` must be of the same length.")
+
     @staticmethod
     def _prepare_video_ids(
         batch_size: int,
@@ -699,7 +724,8 @@ def prepare_latents(
             patch_size=self.transformer_spatial_patch_size,
             device=device,
         )
-        video_ids_scaled = self._scale_video_ids(
+        conditioning_mask = condition_latent_frames_mask.gather(1, video_ids[:, 0])
+        video_ids = self._scale_video_ids(
             video_ids,
             scale_factor=self.vae_spatial_compression_ratio,
             scale_factor_t=self.vae_temporal_compression_ratio,
@@ -709,11 +735,10 @@ def prepare_latents(
         latents = self._pack_latents(
             latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
         )
-        conditioning_mask = condition_latent_frames_mask.gather(1, video_ids[:, 0])
 
         if len(extra_conditioning_latents) > 0:
             latents = torch.cat([*extra_conditioning_latents, latents], dim=1)
-            video_ids = torch.cat([*extra_conditioning_video_ids, video_ids_scaled], dim=2)
+            video_ids = torch.cat([*extra_conditioning_video_ids, video_ids], dim=2)
             conditioning_mask = torch.cat([*extra_conditioning_mask, conditioning_mask], dim=1)
 
         return latents, conditioning_mask, video_ids, extra_conditioning_num_latents
@@ -742,7 +767,11 @@ def interrupt(self):
     @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
-        conditions: Union[LTXVideoCondition, List[LTXVideoCondition]],
+        conditions: Union[LTXVideoCondition, List[LTXVideoCondition]] = None,
+        image: Union[PipelineImageInput, List[PipelineImageInput]] = None,
+        video: List[PipelineImageInput] = None,
+        frame_index: Union[int, List[int]] = 0,
+        strength: Union[float, List[float]] = 1.0,
         prompt: Union[str, List[str]] = None,
         negative_prompt: Optional[Union[str, List[str]]] = None,
         height: int = 512,
@@ -773,8 +802,19 @@ def __call__(
         Function invoked when calling the pipeline for generation.
 
         Args:
-            conditions (`List[LTXVideoCondition]`):
-                The list of frame-conditioning items for the video generation.
+            conditions (`List[LTXVideoCondition], *optional*`):
+                The list of frame-conditioning items for the video generation.If not provided, conditions will be
+                created using `image`, `video`, `frame_index` and `strength`.
+            image (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
+                The image or images to condition the video generation. If not provided, one has to pass `video` or
+                `conditions`.
+            video (`List[PipelineImageInput]`, *optional*):
+                The video to condition the video generation. If not provided, one has to pass `image` or `conditions`.
+            frame_index (`int` or `List[int]`, *optional*):
+                The frame index or frame indices at which the image or video will conditionally effect the video
+                generation. If not provided, one has to pass `conditions`.
+            strength (`float` or `List[float]`, *optional*):
+                The strength or strengths of the conditioning effect. If not provided, one has to pass `conditions`.
             prompt (`str` or `List[str]`, *optional*):
                 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                 instead.
@@ -857,6 +897,11 @@ def __call__(
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
             prompt=prompt,
+            conditions=conditions,
+            image=image,
+            video=video,
+            frame_index=frame_index,
+            strength=strength,
             height=height,
             width=width,
             callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
@@ -878,6 +923,31 @@ def __call__(
         else:
             batch_size = prompt_embeds.shape[0]
 
+        if conditions is not None:
+            if not isinstance(conditions, list):
+                conditions = [conditions]
+
+            strength = [condition.strength for condition in conditions]
+            frame_index = [condition.frame_index for condition in conditions]
+            image = [condition.image for condition in conditions]
+            video = [condition.video for condition in conditions]
+        else:
+            if not isinstance(image, list):
+                image = [image]
+                num_conditions = 1
+            elif isinstance(image, list):
+                num_conditions = len(image)
+            if not isinstance(video, list):
+                video = [video]
+                num_conditions = 1
+            elif isinstance(video, list):
+                num_conditions = len(video)
+
+            if not isinstance(frame_index, list):
+                frame_index = [frame_index] * num_conditions
+            if not isinstance(strength, list):
+                strength = [strength] * num_conditions
+
         device = self._execution_device
 
         # 3. Prepare text embeddings
@@ -905,17 +975,20 @@ def __call__(
         vae_dtype = self.vae.dtype
 
         conditioning_tensors = []
-        conditioning_strengths = []
-        conditioning_start_frames = []
-
-        for condition in conditions:
-            if condition.image is not None:
-                condition_tensor = self.video_processor.preprocess(condition.image, height, width).unsqueeze(2)
-            elif condition.video is not None:
-                condition_tensor = self.video_processor.preprocess_video(condition.video, height, width)
+        for condition_image, condition_video, condition_frame_index, condition_strength in zip(
+            image, video, frame_index, strength
+        ):
+            if condition_image is not None:
+                condition_tensor = (
+                    self.video_processor.preprocess(condition_image, height, width)
+                    .unsqueeze(2)
+                    .to(device, dtype=vae_dtype)
+                )
+            elif condition_video is not None:
+                condition_tensor = self.video_processor.preprocess_video(condition_video, height, width)
                 num_frames_input = condition_tensor.size(2)
                 num_frames_output = self.trim_conditioning_sequence(
-                    condition.frame_index, num_frames_input, num_frames
+                    condition_frame_index, num_frames_input, num_frames
                 )
                 condition_tensor = condition_tensor[:, :, :num_frames_output]
                 condition_tensor = condition_tensor.to(device, dtype=vae_dtype)
@@ -928,15 +1001,13 @@ def __call__(
                     f"but got {condition_tensor.size(2)} frames."
                 )
             conditioning_tensors.append(condition_tensor)
-            conditioning_strengths.append(condition.strength)
-            conditioning_start_frames.append(condition.frame_index)
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels
         latents, conditioning_mask, video_coords, extra_conditioning_num_latents = self.prepare_latents(
             conditioning_tensors,
-            conditioning_strengths,
-            conditioning_start_frames,
+            strength,
+            frame_index,
             batch_size=batch_size * num_videos_per_prompt,
             num_channels_latents=num_channels_latents,
             height=height,
@@ -1015,9 +1086,10 @@ def __call__(
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
                     timestep, _ = timestep.chunk(2)
 
-                denoised_latents = self.scheduler.step(-noise_pred, timestep, latents, return_dict=False)[0]
-                t_eps = 1e-6
-                tokens_to_denoise_mask = (t / 1000 - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
+                denoised_latents = self.scheduler.step(
+                    -noise_pred, t, latents, per_token_timesteps=timestep, return_dict=False
+                )[0]
+                tokens_to_denoise_mask = (t / 1000 - 1e-6 < (1.0 - conditioning_mask)).unsqueeze(-1)
                 latents = torch.where(tokens_to_denoise_mask, denoised_latents, latents)
 
                 if callback_on_step_end is not None: