Merge pull request #87 from Lightricks/pr-2024-12-22

Update ComfyUI-LTXVideo - 2024-12-22

README.md

@@ -1,6 +1,10 @@
 # ComfyUI-LTXVideo
 ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI designed to integrate the LTXVideo diffusion model. These nodes enable workflows for text-to-video, image-to-video, and video-to-video generation. The main LTXVideo repository can be found [here](https://github.com/Lightricks/LTX-Video).
 
+## 22.12.2024
+
+Fixed a bug which caused the model to produce artifacts on short negative prompts when using a native CLIP Loader node.
+
 ## 19.12.2024 ⭐ Update ⭐
 1. Improved model - removes "strobing texture" artifacts and generates better motion. Download from [here](https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.1.safetensors).
 2. STG support

loader_node.py

@@ -74,7 +74,7 @@ class LTXVLoader:
         return (model, vae)
 
     def _load_vae(self, weights, config=None):
-        if config is None:
+        if not config:
             config = {
                 "_class_name": "CausalVideoAutoencoder",
                 "dims": 3,
@@ -121,7 +121,7 @@ class LTXVLoader:
         dtype,
         config=None,
     ):
-        if config is None:
+        if not config:
             config = {
                 "_class_name": "Transformer3DModel",
                 "_diffusers_version": "0.25.1",

model.py

@@ -98,6 +98,13 @@ class LTXVModel(comfy.model_base.BaseModel):
         super().__init__(*args, **kwargs)
         self.model_sampling = LTXVSampling(torch.zeros([1]))
 
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        attention_mask = kwargs.get("attention_mask", None)
+        if attention_mask is not None:
+            out["attention_mask"] = comfy.conds.CONDRegular(attention_mask)
+        return out
+
 
 class LTXVTransformer3D(nn.Module):
     def __init__(
@@ -145,6 +152,7 @@ class LTXVTransformer3D(nn.Module):
         latent,
         timesteps,
         context,
+        attention_mask,
         indices_grid,
         skip_layer_mask=None,
         skip_layer_strategy=None,
@@ -156,7 +164,6 @@ class LTXVTransformer3D(nn.Module):
         # infer mask from context padding, assumes padding vectors are all zero.
         latent = latent.to(self.transformer.dtype)
         latent_patchified = self.patchifier.patchify(latent)
-        context_mask = (context != 0).any(dim=2).to(self.transformer.dtype)
 
         if mixed_precision:
             context_manager = torch.autocast("cuda", dtype=torch.bfloat16)
@@ -169,9 +176,7 @@ class LTXVTransformer3D(nn.Module):
                 ),
                 indices_grid.to(self.transformer.device),
                 encoder_hidden_states=context.to(self.transformer.device),
-                encoder_attention_mask=context_mask.to(self.transformer.device).to(
-                    torch.int64
-                ),
+                encoder_attention_mask=attention_mask,
                 timestep=timesteps,
                 skip_layer_mask=skip_layer_mask,
                 skip_layer_strategy=skip_layer_strategy,
@@ -187,7 +192,16 @@ class LTXVTransformer3D(nn.Module):
         )
         return result
 
-    def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, **kwargs):
+    def forward(
+        self,
+        x,
+        timesteps,
+        context,
+        attention_mask,
+        img_hw=None,
+        aspect_ratio=None,
+        **kwargs,
+    ):
         transformer_options = kwargs.get("transformer_options", {})
         ptb_index = transformer_options.get("ptb_index", None)
         mixed_precision = transformer_options.get("mixed_precision", False)
@@ -211,6 +225,7 @@ class LTXVTransformer3D(nn.Module):
             x,
             timesteps_masked,
             context,
+            attention_mask,
             indices_grid=self.indices_grid(x.shape, x.device),
             mixed_precision=mixed_precision,
             skip_layer_mask=skip_layer_mask,