update

Author: a-r-r-o-w

scripts/convert_ltx_to_diffusers.py

@@ -70,8 +70,6 @@ def remove_keys_(key: str, state_dict: Dict[str, Any]):
     "up_blocks.7": "up_blocks.3.upsamplers.0",
     "up_blocks.8": "up_blocks.3",
     # common
-    "per_channel_scale1": "scale1",
-    "per_channel_scale2": "scale2",
     "last_time_embedder": "time_embedder",
     "last_scale_shift_table": "scale_shift_table",
 }
@@ -168,7 +166,7 @@ def get_vae_config(version: str) -> Dict[str, Any]:
             "decoder_layers_per_block": (4, 3, 3, 3, 4),
             "spatio_temporal_scaling": (True, True, True, False),
             "decoder_spatio_temporal_scaling": (True, True, True, False),
-            "decoder_inject_noise": (False, False, False, False),
+            "decoder_inject_noise": (False, False, False, False, False),
             "upsample_residual": (False, False, False, False),
             "upsample_factor": (1, 1, 1, 1),
             "patch_size": 4,
@@ -190,7 +188,7 @@ def get_vae_config(version: str) -> Dict[str, Any]:
             "decoder_layers_per_block": (5, 6, 7, 8),
             "spatio_temporal_scaling": (True, True, True, False),
             "decoder_spatio_temporal_scaling": (True, True, True),
-            "decoder_inject_noise": (False, True, True, True),
+            "decoder_inject_noise": (True, True, True, False),
             "upsample_residual": (True, True, True),
             "upsample_factor": (2, 2, 2),
             "timestep_conditioning": True,

src/diffusers/models/autoencoders/autoencoder_kl_ltx.py

@@ -138,43 +138,53 @@ def __init__(
                 in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, is_causal=is_causal
             )
 
-        self.scale1 = None
-        self.scale2 = None
+        self.per_channel_scale1 = None
+        self.per_channel_scale2 = None
         if inject_noise:
-            self.scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1))
-            self.scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1))
+            self.per_channel_scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1))
+            self.per_channel_scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1))
 
         self.scale_shift_table = None
         if timestep_conditioning:
             self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
 
-    def forward(self, inputs: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(
+        self, inputs: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None
+    ) -> torch.Tensor:
         hidden_states = inputs
 
         hidden_states = self.norm1(hidden_states.movedim(1, -1)).movedim(-1, 1)
-        scale_1, shift_1, scale_2, shift_2 = self.scale_shift_table.unbind(dim=0)
+
+        if self.scale_shift_table is not None:
+            temb = temb.unflatten(1, (4, -1)) + self.scale_shift_table[None, ..., None, None, None]
+            shift_1, scale_1, shift_2, scale_2 = temb.unbind(dim=1)
+            hidden_states = hidden_states * (1 + scale_1) + shift_1
 
         hidden_states = self.nonlinearity(hidden_states)
         hidden_states = self.conv1(hidden_states)
 
-        if self.scale1 is not None:
+        if self.per_channel_scale1 is not None:
             spatial_shape = hidden_states.shape[-2:]
-            spatial_noise = torch.randn(spatial_shape, device=hidden_states.device, dtype=hidden_states.dtype)
-            hidden_states = hidden_states + (spatial_noise * self.scale1)[None, :, None, :, :]
+            spatial_noise = torch.randn(
+                spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype
+            )
+            hidden_states = hidden_states + (spatial_noise * self.per_channel_scale1)[None, :, None, :, :]
 
         hidden_states = self.norm2(hidden_states.movedim(1, -1)).movedim(-1, 1)
 
         if self.scale_shift_table is not None:
-            hidden_states = hidden_states * (1 + scale_1) + shift_1
+            hidden_states = hidden_states * (1 + scale_2) + shift_2
 
         hidden_states = self.nonlinearity(hidden_states)
         hidden_states = self.dropout(hidden_states)
         hidden_states = self.conv2(hidden_states)
 
-        if self.scale2 is not None:
+        if self.per_channel_scale2 is not None:
             spatial_shape = hidden_states.shape[-2:]
-            spatial_noise = torch.randn(spatial_shape, device=hidden_states.device, dtype=hidden_states.dtype)
-            hidden_states = hidden_states + (spatial_noise * self.scale2)[None, :, None, :, :]
+            spatial_noise = torch.randn(
+                spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype
+            )
+            hidden_states = hidden_states + (spatial_noise * self.per_channel_scale2)[None, :, None, :, :]
 
         if self.norm3 is not None:
             inputs = self.norm3(inputs.movedim(1, -1)).movedim(-1, 1)
@@ -318,7 +328,12 @@ def __init__(
 
         self.gradient_checkpointing = False
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
         r"""Forward method of the `LTXDownBlock3D` class."""
 
         for i, resnet in enumerate(self.resnets):
@@ -330,16 +345,18 @@ def create_forward(*inputs):
 
                     return create_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb, generator
+                )
             else:
-                hidden_states = resnet(hidden_states)
+                hidden_states = resnet(hidden_states, temb, generator)
 
         if self.downsamplers is not None:
             for downsampler in self.downsamplers:
                 hidden_states = downsampler(hidden_states)
 
         if self.conv_out is not None:
-            hidden_states = self.conv_out(hidden_states)
+            hidden_states = self.conv_out(hidden_states, temb, generator)
 
         return hidden_states
 
@@ -401,7 +418,12 @@ def __init__(
 
         self.gradient_checkpointing = False
 
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
         r"""Forward method of the `LTXMidBlock3D` class."""
 
         if self.time_embedder is not None:
@@ -423,9 +445,11 @@ def create_forward(*inputs):
 
                     return create_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb, generator
+                )
             else:
-                hidden_states = resnet(hidden_states, temb)
+                hidden_states = resnet(hidden_states, temb, generator)
 
         return hidden_states
 
@@ -524,9 +548,14 @@ def __init__(
 
         self.gradient_checkpointing = False
 
-    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
         if self.conv_in is not None:
-            hidden_states = self.conv_in(hidden_states)
+            hidden_states = self.conv_in(hidden_states, temb, generator)
 
         if self.time_embedder is not None:
             temb = self.time_embedder(
@@ -551,9 +580,11 @@ def create_forward(*inputs):
 
                     return create_forward
 
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb, generator
+                )
             else:
-                hidden_states = resnet(hidden_states)
+                hidden_states = resnet(hidden_states, temb, generator)
 
         return hidden_states
 
@@ -746,6 +777,9 @@ def __init__(
         block_out_channels = tuple(reversed(block_out_channels))
         spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
         layers_per_block = tuple(reversed(layers_per_block))
+        inject_noise = tuple(reversed(inject_noise))
+        upsample_residual = tuple(reversed(upsample_residual))
+        upsample_factor = tuple(reversed(upsample_factor))
         output_channel = block_out_channels[0]
 
         self.conv_in = LTXCausalConv3d(
@@ -810,29 +844,31 @@ def create_forward(*inputs):
 
                 return create_forward
 
-            hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), hidden_states)
+            hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self.mid_block), hidden_states, temb
+            )
 
             for up_block in self.up_blocks:
-                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), hidden_states)
+                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), hidden_states, temb)
         else:
-            hidden_states = self.mid_block(hidden_states)
+            hidden_states = self.mid_block(hidden_states, temb)
 
             for up_block in self.up_blocks:
-                hidden_states = up_block(hidden_states)
+                hidden_states = up_block(hidden_states, temb)
 
         hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
 
         if self.time_embedder is not None:
-            embedded_timestep = self.time_embedder(
+            temb = self.time_embedder(
                 timestep=temb.flatten(),
                 resolution=None,
                 aspect_ratio=None,
                 batch_size=hidden_states.size(0),
                 hidden_dtype=hidden_states.dtype,
             )
-            embedded_timestep = embedded_timestep.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1))
-            embedded_timestep = embedded_timestep + self.scale_shift_table[None, :, None, None, None]
-            shift, scale = embedded_timestep.unbind(dim=1)
+            temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1))
+            temb = temb + self.scale_shift_table[None, ..., None, None, None]
+            shift, scale = temb.unbind(dim=1)
             hidden_states = hidden_states * (1 + scale) + shift
 
         hidden_states = self.conv_act(hidden_states)
@@ -902,7 +938,7 @@ def __init__(
         decoder_layers_per_block: Tuple[int, ...] = (4, 3, 3, 3, 4),
         spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
         decoder_spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, False),
-        decoder_inject_noise: Tuple[bool, ...] = (False, False, False, False),
+        decoder_inject_noise: Tuple[bool, ...] = (False, False, False, False, False),
         upsample_residual: Tuple[bool, ...] = (False, False, False, False),
         upsample_factor: Tuple[int, ...] = (1, 1, 1, 1),
         timestep_conditioning: bool = False,
@@ -1078,13 +1114,15 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def _decode(
+        self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio
 
         if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
-            return self.tiled_decode(z, return_dict=return_dict)
+            return self.tiled_decode(z, temb, return_dict=return_dict)
 
         if self.use_framewise_decoding:
             # TODO(aryan): requires investigation
@@ -1094,15 +1132,17 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
                 "should be possible, please submit a PR to https://github.com/huggingface/diffusers/pulls."
             )
         else:
-            dec = self.decoder(z)
+            dec = self.decoder(z, temb)
 
         if not return_dict:
             return (dec,)
 
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def decode(
+        self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -1117,10 +1157,15 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 returned.
         """
         if self.use_slicing and z.shape[0] > 1:
-            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            if temb is not None:
+                decoded_slices = [
+                    self._decode(z_slice, t_slice).sample for z_slice, t_slice in (z.split(1), temb.split(1))
+                ]
+            else:
+                decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
             decoded = torch.cat(decoded_slices)
         else:
-            decoded = self._decode(z).sample
+            decoded = self._decode(z, temb).sample
 
         if not return_dict:
             return (decoded,)
@@ -1202,7 +1247,9 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
         enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def tiled_decode(
+        self, z: torch.Tensor, temb: Optional[torch.Tensor], return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1243,7 +1290,9 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[Decod
                         "should be possible, please submit a PR to https://github.com/huggingface/diffusers/pulls."
                     )
                 else:
-                    time = self.decoder(z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width])
+                    time = self.decoder(
+                        z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width], temb
+                    )
 
                 row.append(time)
             rows.append(row)
@@ -1271,6 +1320,7 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[Decod
     def forward(
         self,
         sample: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
@@ -1281,7 +1331,7 @@ def forward(
             z = posterior.sample(generator=generator)
         else:
             z = posterior.mode()
-        dec = self.decode(z)
+        dec = self.decode(z, temb)
         if not return_dict:
             return (dec,)
         return dec

src/diffusers/pipelines/ltx/pipeline_ltx.py

@@ -511,6 +511,8 @@ def __call__(
         prompt_attention_mask: Optional[torch.Tensor] = None,
         negative_prompt_embeds: Optional[torch.Tensor] = None,
         negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        decode_timestep: Union[float, List[float]] = 0.05,
+        decode_noise_scale: Union[float, List[float]] = 0.025,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
         attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -753,7 +755,25 @@ def __call__(
                 latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
             )
             latents = latents.to(prompt_embeds.dtype)
-            video = self.vae.decode(latents, return_dict=False)[0]
+
+            if not self.vae.config.timestep_conditioning:
+                timestep = None
+            else:
+                noise = torch.randn(latents.shape, generator=generator, device=device, dtype=latents.dtype)
+                if not isinstance(decode_timestep, list):
+                    decode_timestep = [decode_timestep] * batch_size
+                if decode_noise_scale is None:
+                    decode_noise_scale = decode_timestep
+                elif not isinstance(decode_noise_scale, list):
+                    decode_noise_scale = [decode_noise_scale] * batch_size
+
+                timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
+                decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
+                    :, None, None, None, None
+                ]
+                latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
+
+            video = self.vae.decode(latents, timestep, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)
 
         # Offload all models