Automated PR - 2026-01-27

Update ComfyUI-LTXVideo - 2026-01-07

__init__.py

@@ -5,7 +5,9 @@ from .easy_samplers import (
     LTXVBaseSampler,
     LTXVExtendSampler,
     LTXVInContextSampler,
+    LTXVNormalizingSampler,
 )
+from .gemma_api_conditioning import GemmaAPITextEncode
 from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
 from .guide import LTXVAddGuideAdvanced
 from .guiders import GuiderParametersNode, MultimodalGuiderNode
@@ -61,6 +63,7 @@ NODE_CLASS_MAPPINGS = {
     "LTXVBaseSampler": LTXVBaseSampler,
     "LTXVInContextSampler": LTXVInContextSampler,
     "LTXVExtendSampler": LTXVExtendSampler,
+    "LTXVNormalizingSampler": LTXVNormalizingSampler,
     "LTXVPreprocessMasks": LTXVPreprocessMasks,
     "LTXVPatcherVAE": LTXVPatcherVAE,
     "LTXVPromptEnhancer": LTXVPromptEnhancer,
@@ -83,6 +86,7 @@ NODE_CLASS_MAPPINGS = {
     "LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
     "LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
     "LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
+    "GemmaAPITextEncode": GemmaAPITextEncode,
     "DynamicConditioning": DynamicConditioning,
     "LowVRAMCheckpointLoader": LowVRAMCheckpointLoader,
     "LowVRAMAudioVAELoader": LowVRAMAudioVAELoader,

easy_samplers.py

@@ -4,6 +4,9 @@ import comfy
 import comfy_extras
 import nodes
 import torch
+from comfy.ldm.lightricks.av_model import LTXAVModel
+from comfy.nested_tensor import NestedTensor
+from comfy_api.latest import io
 from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced, SplitSigmas
 from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGuides
 
@@ -907,3 +910,142 @@ class LinearOverlapLatentTransition:
                 "batch_index": combined_batch_index,
             },
         )
+
+
+@comfy_node(description="LTXV Normalizing Sampler")
+class LTXVNormalizingSampler(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVNormalizingSampler",
+            category="utility",
+            inputs=[
+                io.Noise.Input("noise"),
+                io.Guider.Input("guider"),
+                io.Sampler.Input("sampler"),
+                io.Sigmas.Input("sigmas"),
+                io.Latent.Input("latent_image"),
+                io.String.Input(
+                    "video_normalization_factors", default="1,1,1,1,1,1,1,1"
+                ),
+                io.String.Input(
+                    "audio_normalization_factors", default="1,1,0.25,1,1,0.25,1,1"
+                ),
+            ],
+            outputs=[
+                io.Latent.Output(display_name="denoised_output"),
+            ],
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        noise,
+        guider,
+        sampler,
+        sigmas,
+        latent_image,
+        video_normalization_factors,
+        audio_normalization_factors,
+    ) -> io.NodeOutput:
+
+        if (
+            guider.model_patcher.model.diffusion_model.__class__.__name__
+            != "LTXAVModel"
+        ):
+            raise ValueError()
+
+        ltxav: LTXAVModel = guider.model_patcher.model.diffusion_model
+
+        video_normalization_factors = video_normalization_factors.split(",")
+        audio_normalization_factors = audio_normalization_factors.split(",")
+        video_normalization_factors = [
+            float(factor) for factor in video_normalization_factors
+        ]
+        audio_normalization_factors = [
+            float(factor) for factor in audio_normalization_factors
+        ]
+
+        # Extend normalization factors to match the length of sigmas
+        sigmas_len = len(sigmas) - 1
+        if len(video_normalization_factors) < sigmas_len:
+            if len(video_normalization_factors) > 0:
+                video_normalization_factors.extend(
+                    [video_normalization_factors[-1]]
+                    * (sigmas_len - len(video_normalization_factors))
+                )
+        if len(audio_normalization_factors) < sigmas_len:
+            if len(audio_normalization_factors) > 0:
+                audio_normalization_factors.extend(
+                    [audio_normalization_factors[-1]]
+                    * (sigmas_len - len(audio_normalization_factors))
+                )
+
+        # Calculate indices where both normalization factors are not 1.0
+        sampling_split_indices = [
+            i + 1
+            for i, (v, a) in enumerate(
+                zip(video_normalization_factors, audio_normalization_factors)
+            )
+            if v != 1.0 or a != 1.0
+        ]
+        print("Sampling split indices: %s" % sampling_split_indices, flush=True)
+
+        # Split sigmas according to sampling_split_indices
+        def split_by_indices(arr, indices):
+            """
+            Splits arr into chunks according to indices (split points).
+            Indices are treated as starting a new chunk at each index in the list.
+            """
+            if not indices:
+                return [arr]
+            split_points = sorted(set(indices))
+            chunks = []
+            prev = 0
+            for idx in split_points:
+                if prev < idx:
+                    chunks.append(arr[prev : idx + 1])
+                prev = idx
+            if prev < len(arr):
+                chunks.append(arr[prev:])
+            return chunks
+
+        sigmas_chunks = split_by_indices(sigmas, sampling_split_indices)
+        print("Sigmas chunks: %s" % sigmas_chunks, flush=True)
+
+        i = 0
+        for sigmas_chunk in sigmas_chunks:
+            i += len(sigmas_chunk) - 1
+            print("Sampling with sigmas %s" % (sigmas_chunk), flush=True)
+            (_, latent_image) = SamplerCustomAdvanced().execute(
+                noise=noise,
+                guider=guider,
+                sampler=sampler,
+                sigmas=sigmas_chunk,
+                latent_image=latent_image,
+            )
+            video_samples, audio_samples = ltxav.separate_audio_and_video_latents(
+                latent_image["samples"].tensors,
+                None,
+            )
+            if i - 1 < len(video_normalization_factors) and i - 1 < len(
+                audio_normalization_factors
+            ):
+                video_samples = video_samples * video_normalization_factors[i - 1]
+                audio_samples = audio_samples * audio_normalization_factors[i - 1]
+                latent_image["samples"] = NestedTensor(
+                    ltxav.recombine_audio_and_video_latents(
+                        video_samples, audio_samples
+                    )
+                )
+                print(
+                    "After %d steps, the latent image was normalized by %f and %f"
+                    % (
+                        i,
+                        video_normalization_factors[i - 1],
+                        audio_normalization_factors[i - 1],
+                    ),
+                    flush=True,
+                )
+
+        return io.NodeOutput(latent_image)

example_workflows/LTX-2_I2V_Distilled_wLora.json

@@ -565,7 +565,7 @@
       "title": "Video Size",
       "properties": {},
       "widgets_values": [
-        "Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
+        "Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
       ],
       "color": "#432",
       "bgcolor": "#653"

example_workflows/LTX-2_I2V_Full_wLora.json

@@ -743,7 +743,7 @@
       "title": "Video Size",
       "properties": {},
       "widgets_values": [
-        "Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
+        "Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
       ],
       "color": "#432",
       "bgcolor": "#653"

example_workflows/LTX-2_ICLoRA_All_Distilled.json

@@ -57,7 +57,7 @@
       "title": "Video Size",
       "properties": {},
       "widgets_values": [
-        "Width & height settings must be divisible by 32 + 1. \n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
+        "Width & height settings must be divisible by 64. \n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
       ],
       "color": "#432",
       "bgcolor": "#653"

example_workflows/LTX-2_T2V_Distilled_wLora.json

@@ -567,7 +567,7 @@
       "title": "Video Size",
       "properties": {},
       "widgets_values": [
-        "Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
+        "Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
       ],
       "color": "#432",
       "bgcolor": "#653"

example_workflows/LTX-2_T2V_Full_wLora.json

@@ -520,7 +520,7 @@
       "title": "Video Size",
       "properties": {},
       "widgets_values": [
-        "Width & height settings must be divisible by 32 + 1. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
+        "Width & height settings must be divisible by 64. \nFrame count must be divisible by 8 + 1. \n\n\nRunning with invalid parameters **will not cause errors**. Instead, the flow will silently choose the closest valid parameters. "
       ],
       "color": "#432",
       "bgcolor": "#653"

gemma_api_conditioning.py

@@ -0,0 +1,134 @@
+"""
+API-based text encoding that returns CONDITIONING for LTX-2.
+Replaces the CLIP encoding step entirely using an external API.
+"""
+
+import io
+import logging
+import pickle
+
+import folder_paths
+import requests
+from safetensors import safe_open
+
+from .nodes_registry import comfy_node
+
+logger = logging.getLogger(__name__)
+
+LTXV_API_BASE_URL = "https://api.ltx.video"
+UPDATE_MESSAGE = (
+    "Note: If this error persists, the node might be outdated. "
+    "Please update ComfyUI-LTXVideo to the latest version."
+)
+INVALID_API_KEY_MESSAGE = (
+    "Invalid API key. Please generate a new API key at: https://console.ltx.video/"
+)
+MISSING_MODEL_ID_MESSAGE = "Model ID cannot be identified from the provided model file"
+
+
+def extract_model_id(ckpt_name: str) -> str:
+    model_id_key = "encrypted_wandb_properties"
+    with safe_open(
+        folder_paths.get_full_path_or_raise("checkpoints", ckpt_name),
+        framework="pt",
+        device="cpu",
+    ) as f:
+        metadata = f.metadata()
+        if not metadata or model_id_key not in metadata:
+            raise ValueError(MISSING_MODEL_ID_MESSAGE)
+        return metadata[model_id_key]
+
+
+@comfy_node(name="GemmaAPITextEncode")
+class GemmaAPITextEncode:
+    """
+    Encodes text prompts using the LTX Video API, returning CONDITIONING for LTX-2 models.
+
+    This node replaces the local CLIP encoding step by sending the prompt to an external API
+    for processing. It requires an API key and automatically extracts the model ID from the
+    checkpoint file metadata.
+
+    Inputs:
+        - api_key: Authentication key for the LTX Video API
+        - prompt: Text prompt to encode
+        - ckpt_name: Checkpoint file containing model metadata
+
+    Returns:
+        - CONDITIONING: Encoded prompt conditioning ready for LTX-2 video generation
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "api_key": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "placeholder": "API_KEY",
+                        "multiline": False,
+                        "tooltip": "API key for authentication",
+                    },
+                ),
+                "prompt": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Text prompt to encode",
+                    },
+                ),
+                "ckpt_name": (
+                    folder_paths.get_filename_list("checkpoints"),
+                    {"tooltip": "The name of the checkpoint (model) to load."},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    RETURN_NAMES = ("conditioning",)
+    FUNCTION = "encode"
+    CATEGORY = "api node/text/Lightricks"
+
+    def encode(self, api_key: str, prompt: str, ckpt_name: str):
+        if not api_key:
+            raise ValueError("API key is required")
+
+        if not prompt.strip():
+            raise ValueError("Text prompt cannot be empty")
+
+        if not ckpt_name or not ckpt_name.strip():
+            raise ValueError("Model path is required")
+
+        model_id = extract_model_id(ckpt_name)
+        payload = {"prompt": prompt, "model_id": model_id}
+        logger.info(
+            f"Calling API to encode prompt: {prompt[:50]}... with model_id: {model_id[:50]}..."
+        )
+        try:
+            response = requests.post(
+                f"{LTXV_API_BASE_URL}/v1/prompt-embedding",
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {api_key}",
+                    "Content-Type": "application/json",
+                },
+                timeout=60,
+            )
+
+            if response.status_code == 401:
+                raise RuntimeError(INVALID_API_KEY_MESSAGE)
+
+            if response.status_code != 200:
+                raise RuntimeError(
+                    f"API request failed with status {response.status_code}: {response.text}\n"
+                    f"{UPDATE_MESSAGE}"
+                )
+
+            conditioning = pickle.load(io.BytesIO(response.content))
+            logger.info("Successfully received conditioning from API")
+
+            return (conditioning,)
+
+        except Exception as e:
+            raise RuntimeError(f"API request failed: {str(e)}\n {UPDATE_MESSAGE}")

gemma_configs/gemma3cfg.json

@@ -0,0 +1,62 @@
+{
+    "architectures": [
+      "Gemma3ForConditionalGeneration"
+    ],
+    "boi_token_index": 255999,
+    "eoi_token_index": 256000,
+    "eos_token_id": [
+      1,
+      106
+    ],
+    "image_token_index": 262144,
+    "initializer_range": 0.02,
+    "mm_tokens_per_image": 256,
+    "model_type": "gemma3",
+    "text_config": {
+      "attention_bias": false,
+      "attention_dropout": 0.0,
+      "attn_logit_softcapping": null,
+      "cache_implementation": "hybrid",
+      "final_logit_softcapping": null,
+      "head_dim": 256,
+      "hidden_activation": "gelu_pytorch_tanh",
+      "hidden_size": 3840,
+      "initializer_range": 0.02,
+      "intermediate_size": 15360,
+      "max_position_embeddings": 131072,
+      "model_type": "gemma3_text",
+      "num_attention_heads": 16,
+      "num_hidden_layers": 48,
+      "num_key_value_heads": 8,
+      "query_pre_attn_scalar": 256,
+      "rms_norm_eps": 1e-06,
+      "rope_local_base_freq": 10000,
+      "rope_scaling": {
+        "factor": 8.0,
+        "rope_type": "linear"
+      },
+      "rope_theta": 1000000,
+      "sliding_window": 1024,
+      "sliding_window_pattern": 6,
+      "torch_dtype": "bfloat16",
+      "use_cache": true,
+      "vocab_size": 262208
+    },
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.51.0",
+    "vision_config": {
+      "attention_dropout": 0.0,
+      "hidden_act": "gelu_pytorch_tanh",
+      "hidden_size": 1152,
+      "image_size": 896,
+      "intermediate_size": 4304,
+      "layer_norm_eps": 1e-06,
+      "model_type": "siglip_vision_model",
+      "num_attention_heads": 16,
+      "num_channels": 3,
+      "num_hidden_layers": 27,
+      "patch_size": 14,
+      "torch_dtype": "bfloat16",
+      "vision_use_head": false
+    }
+  }

gemma_configs/preprocessor_config.json

@@ -0,0 +1,29 @@
+{
+  "do_convert_rgb": null,
+  "do_normalize": true,
+  "do_pan_and_scan": null,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Gemma3ImageProcessor",
+  "image_seq_length": 256,
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "pan_and_scan_max_num_crops": null,
+  "pan_and_scan_min_crop_size": null,
+  "pan_and_scan_min_ratio_to_activate": null,
+  "processor_class": "Gemma3Processor",
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 896,
+    "width": 896
+  }
+}

gemma_configs/processor_config.json

@@ -0,0 +1,4 @@
+{
+  "image_seq_length": 256,
+  "processor_class": "Gemma3Processor"
+}

gemma_encoder.py

@@ -1,7 +1,8 @@
 import json
 import logging
+from glob import glob
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional, Tuple
 
 import comfy.model_management
 import comfy.ops
@@ -17,9 +18,12 @@ from PIL import Image
 from transformers import (
     AutoImageProcessor,
     AutoTokenizer,
+    Gemma3Config,
     Gemma3ForConditionalGeneration,
     Gemma3Processor,
 )
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.models.auto.modeling_auto import MODEL_MAPPING_NAMES
 
 from .embeddings_connector import Embeddings1DConnector
 from .nodes_registry import comfy_node
@@ -351,136 +355,6 @@ class LTXVGemmaTextEncoderModel(torch.nn.Module):
         # Return a conservative estimate in bytesed(input_shape)
         return self._model_memory_required
 
-    @staticmethod
-    def _cat_with_padding(
-        tensor: torch.Tensor,
-        padding_length: int,
-        value: int | float,
-    ) -> torch.Tensor:
-        """Concatenate a tensor with a padding tensor of the given value."""
-        return torch.cat(
-            [
-                tensor,
-                torch.full(
-                    (1, padding_length),
-                    value,
-                    dtype=tensor.dtype,
-                    device=tensor.device,
-                ),
-            ],
-            dim=1,
-        )
-
-    def _pad_inputs_for_attention_alignment(self, model_inputs, alignment: int = 8):
-        """Pad sequence length to multiple of alignment for Flash Attention compatibility.
-
-        Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
-        This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
-        'p.attn_bias_ptr is not correctly aligned' errors.
-        """
-        seq_len = model_inputs.input_ids.shape[1]
-        padded_len = ((seq_len + alignment - 1) // alignment) * alignment
-        padding_length = padded_len - seq_len
-
-        if padding_length > 0:
-            pad_token_id = (
-                self.processor.tokenizer.pad_token_id
-                if self.processor.tokenizer.pad_token_id is not None
-                else 0
-            )
-
-            model_inputs["input_ids"] = self._cat_with_padding(
-                model_inputs.input_ids, padding_length, pad_token_id
-            )
-
-            model_inputs["attention_mask"] = self._cat_with_padding(
-                model_inputs.attention_mask, padding_length, 0
-            )
-
-            if (
-                "token_type_ids" in model_inputs
-                and model_inputs["token_type_ids"] is not None
-            ):
-                model_inputs["token_type_ids"] = self._cat_with_padding(
-                    model_inputs["token_type_ids"], padding_length, 0
-                )
-
-        return model_inputs
-
-    def _enhance(
-        self,
-        messages: list,
-        image: Optional[Image.Image] = None,
-        max_new_tokens: int = 512,
-        seed: int = 42,
-    ) -> str:
-        """Common enhancement logic for both T2V and I2V modes."""
-        if self.processor is None:
-            raise ValueError("Processor not loaded - enhancement not available")
-
-        text = self.processor.tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-
-        model_inputs = self.processor(
-            text=text,
-            images=image,
-            return_tensors="pt",
-        ).to(self.model.device)
-        model_inputs = self._pad_inputs_for_attention_alignment(model_inputs)
-
-        with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
-            torch.manual_seed(seed)
-            outputs = self.model.generate(
-                **model_inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=True,
-                temperature=0.7,
-            )
-            generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
-            enhanced_prompt = self.processor.tokenizer.decode(
-                generated_ids, skip_special_tokens=True
-            )
-
-        return enhanced_prompt
-
-    def enhance_t2v(
-        self,
-        prompt: str,
-        system_prompt: str,
-        max_new_tokens: int = 512,
-        seed: int = 42,
-    ) -> str:
-        """Enhance a text prompt for T2V generation."""
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
-        ]
-        return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
-
-    def enhance_i2v(
-        self,
-        prompt: str,
-        image: Image.Image,
-        system_prompt: str,
-        max_new_tokens: int = 512,
-        seed: int = 42,
-    ) -> str:
-        """Enhance a text prompt for I2V generation using a reference image."""
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
-                ],
-            },
-        ]
-        return self._enhance(
-            messages, image=image, max_new_tokens=max_new_tokens, seed=seed
-        )
-
 
 def ltxv_gemma_tokenizer(tokenizer_path, max_length=256):
     class _LTXVGemmaTokenizer(LTXVGemmaTokenizer):
@@ -528,8 +402,11 @@ def find_matching_dir(root_path: str, pattern: str) -> str:
     """
     Recursively search for files matching a glob pattern and return the parent directory of the first match.
     """
-
-    matches = list(Path(root_path).rglob(pattern))
+    matches = [
+        Path(p)
+        for p in glob(f"{root_path}/**", recursive=True)
+        if Path(p).match(pattern)
+    ]
     if not matches:
         raise FileNotFoundError(
             f"No files matching pattern '{pattern}' found under {root_path}"
@@ -672,13 +549,18 @@ class LTXVGemmaEnhancePrompt:
         clip.load_model()
         encoder = clip.cond_stage_model
 
-        if encoder.processor is None:
-            raise ValueError(
-                "Processor not loaded - enhancement not available. "
-                "Ensure your model directory has chat_template.json, processor_config.json, "
-                "and preprocessor_config.json files."
-            )
-
+        if not hasattr(encoder, "processor") or encoder.processor is None:
+            if hasattr(encoder, "gemma3_12b"):
+                model, processor = transformers_gemma3_from_encoder(encoder)
+            else:
+                raise ValueError(
+                    "Processor not loaded - enhancement not available. "
+                    "Ensure your model directory has chat_template.json, processor_config.json, "
+                    "and preprocessor_config.json files."
+                )
+        elif isinstance(encoder, LTXVGemmaTextEncoderModel):
+            model = encoder.model
+            processor = encoder.processor
         # Determine mode: use I2V if image is provided and not bypassed
         use_i2v = image is not None and not bypass_i2v
 
@@ -694,7 +576,9 @@ class LTXVGemmaEnhancePrompt:
 
         if use_i2v:
             pil_image = tensor_to_pil(image)
-            enhanced_prompt = encoder.enhance_i2v(
+            enhanced_prompt = enhance_i2v(
+                processor=processor,
+                model=model,
                 prompt=prompt,
                 image=pil_image,
                 system_prompt=system_prompt,
@@ -702,7 +586,9 @@ class LTXVGemmaEnhancePrompt:
                 seed=seed,
             )
         else:
-            enhanced_prompt = encoder.enhance_t2v(
+            enhanced_prompt = enhance_t2v(
+                processor=processor,
+                model=model,
                 prompt=prompt,
                 system_prompt=system_prompt,
                 max_new_tokens=max_tokens,
@@ -712,3 +598,275 @@ class LTXVGemmaEnhancePrompt:
         enhanced_prompt = clean_response(enhanced_prompt)
 
         return (enhanced_prompt,)
+
+
+def _enhance(
+    processor: Gemma3Processor,
+    model: Gemma3ForConditionalGeneration,
+    messages: list,
+    image: Optional[Image.Image] = None,
+    max_new_tokens: int = 512,
+    seed: int = 42,
+) -> str:
+    """Common enhancement logic for both T2V and I2V modes."""
+    if processor is None:
+        raise ValueError("Processor not loaded - enhancement not available")
+
+    text = processor.tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    model_inputs = processor(
+        text=text,
+        images=image,
+        return_tensors="pt",
+    ).to(model.device)
+
+    pad_token_id = (
+        processor.tokenizer.pad_token_id
+        if processor.tokenizer.pad_token_id is not None
+        else 0
+    )
+    model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id)
+
+    with (
+        torch.inference_mode(),
+        torch.random.fork_rng(devices=[model.device]),
+        torch.autocast(device_type=model.device.type, dtype=model.dtype),
+    ):
+        torch.manual_seed(seed)
+        outputs = model.generate(
+            **model_inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+        )
+        generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
+        enhanced_prompt = processor.tokenizer.decode(
+            generated_ids, skip_special_tokens=True
+        )
+
+    return enhanced_prompt
+
+
+def enhance_t2v(
+    processor: Gemma3Processor,
+    model: Gemma3ForConditionalGeneration,
+    prompt: str,
+    system_prompt: str,
+    max_new_tokens: int = 512,
+    seed: int = 42,
+) -> str:
+    """Enhance a text prompt for T2V generation."""
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
+    ]
+    return _enhance(
+        processor, model, messages, max_new_tokens=max_new_tokens, seed=seed
+    )
+
+
+def enhance_i2v(
+    processor: Gemma3Processor,
+    model: Gemma3ForConditionalGeneration,
+    prompt: str,
+    image: Image.Image,
+    system_prompt: str,
+    max_new_tokens: int = 512,
+    seed: int = 42,
+) -> str:
+    """Enhance a text prompt for I2V generation using a reference image."""
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
+            ],
+        },
+    ]
+    return _enhance(
+        processor,
+        model,
+        messages,
+        image=image,
+        max_new_tokens=max_new_tokens,
+        seed=seed,
+    )
+
+
+def _cat_with_padding(
+    tensor: torch.Tensor,
+    padding_length: int,
+    value: int | float,
+) -> torch.Tensor:
+    """Concatenate a tensor with a padding tensor of the given value."""
+    return torch.cat(
+        [
+            tensor,
+            torch.full(
+                (1, padding_length),
+                value,
+                dtype=tensor.dtype,
+                device=tensor.device,
+            ),
+        ],
+        dim=1,
+    )
+
+
+def _pad_inputs_for_attention_alignment(model_inputs, pad_token_id, alignment: int = 8):
+    """Pad sequence length to multiple of alignment for Flash Attention compatibility.
+
+    Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
+    This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
+    'p.attn_bias_ptr is not correctly aligned' errors.
+    """
+    seq_len = model_inputs.input_ids.shape[1]
+    padded_len = ((seq_len + alignment - 1) // alignment) * alignment
+    padding_length = padded_len - seq_len
+
+    if padding_length > 0:
+        model_inputs["input_ids"] = _cat_with_padding(
+            model_inputs.input_ids, padding_length, pad_token_id
+        )
+
+        model_inputs["attention_mask"] = _cat_with_padding(
+            model_inputs.attention_mask, padding_length, 0
+        )
+
+        if (
+            "token_type_ids" in model_inputs
+            and model_inputs["token_type_ids"] is not None
+        ):
+            model_inputs["token_type_ids"] = _cat_with_padding(
+                model_inputs["token_type_ids"], padding_length, 0
+            )
+
+    return model_inputs
+
+
+def _locate_model_within_model(super_model, model_name):
+    class_name = MODEL_MAPPING_NAMES.get(model_name, None)
+    if class_name is None:
+        return None
+    for module in super_model.modules():
+        if module.__class__.__name__ == class_name:
+            return module
+    return None
+
+
+def _locate_unique_parameter_owner_by_leaf(
+    root: torch.nn.Module,
+    leaf_param_name: str,
+    must_have_in_path: Optional[str] = None,
+) -> Optional[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]]:
+
+    modules = dict(root.named_modules())
+
+    candidates: List[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]] = []
+    for full_name, p in root.named_parameters(recurse=True):
+        parts = full_name.split(".")
+        leaf = parts[-1]
+        if leaf != leaf_param_name:
+            continue
+        if must_have_in_path is not None and must_have_in_path not in parts:
+            continue
+
+        owner_path = ".".join(parts[:-1])
+        owner = modules.get(owner_path, root if owner_path == "" else None)
+        if owner is None:
+            continue
+        candidates.append((owner, leaf, p, full_name))
+
+    if not candidates:
+        return None
+    return candidates[0]
+
+
+def transformers_gemma3_from_encoder(encoder):
+    jsons_path = Path(__file__).parent / "gemma_configs"
+    config = Gemma3Config.from_json_file(jsons_path / "gemma3cfg.json")
+    with torch.device("meta"):
+        metamodel = Gemma3ForConditionalGeneration(config)
+    t_model_name = config.text_config.model_type
+    t_model = _locate_model_within_model(metamodel, t_model_name)
+    if t_model is None:
+        raise ValueError(
+            "Can't locate text model while converting text encoder to Gemma3ForConditionalGeneration"
+        )
+    t_model.load_state_dict(
+        encoder.gemma3_12b.transformer.model.state_dict(), assign=True, strict=False
+    )
+    v_tower_name = config.vision_config.model_type
+    v_tower = _locate_model_within_model(metamodel, v_tower_name)
+    if v_tower is None:
+        raise ValueError(
+            "Can't locate vision model while converting text encoder to Gemma3ForConditionalGeneration"
+        )
+    v_model = v_tower.vision_model
+    v_model.load_state_dict(
+        encoder.gemma3_12b.transformer.vision_model.state_dict(),
+        assign=True,
+        strict=False,
+    )
+    metamodel.multi_modal_projector.load_state_dict(
+        encoder.gemma3_12b.transformer.multi_modal_projector.state_dict(),
+        assign=True,
+        strict=False,
+    )
+    config = config.text_config
+    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    base = config.rope_local_base_freq
+
+    device = encoder.device
+    positions_length = len(v_model.embeddings.position_ids[0])
+    position_ids = torch.arange(
+        positions_length, dtype=torch.long, device="cpu"
+    ).unsqueeze(0)
+    v_model.embeddings.register_buffer("position_ids", position_ids)
+    embed_scale = torch.tensor(config.hidden_size**0.5, device=device)
+    t_model.embed_tokens.register_buffer("embed_scale", embed_scale)
+    local_rope_freqs = 1.0 / (
+        base
+        ** (
+            torch.arange(0, dim, 2, dtype=torch.int64).to(
+                device=device, dtype=torch.float
+            )
+            / dim
+        )
+    )
+    t_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
+    rope_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](
+        config, device
+    )
+    t_model.rotary_emb.register_buffer("inv_freq", rope_freqs)
+    lm_head_requires_grad = False
+    loc_result = _locate_unique_parameter_owner_by_leaf(
+        metamodel, leaf_param_name="weight", must_have_in_path="lm_head"
+    )
+    if loc_result is None:
+        raise ValueError(
+            "Can't locate lm_head while converting text encoder to Gemma3ForConditionalGeneration"
+        )
+    lm_head_owner, lm_head_attr, _, _ = loc_result
+    real_w = t_model.embed_tokens.weight
+    setattr(
+        lm_head_owner,
+        lm_head_attr,
+        torch.nn.Parameter(real_w, requires_grad=lm_head_requires_grad),
+    )
+    metamodel.to(device)
+
+    tokenizer_class = ltxv_gemma_tokenizer(jsons_path, max_length=1024)
+    image_processor = AutoImageProcessor.from_pretrained(
+        str(jsons_path),
+        local_files_only=True,
+    )
+    processor = Gemma3Processor(
+        image_processor=image_processor,
+        tokenizer=tokenizer_class().tokenizer,
+    )
+    return metamodel, processor

guiders/multimodal_guider.py

@@ -64,7 +64,6 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
         self,
         run_vx: bool,
         run_ax: bool,
-        audio_length: int,
         audio_ptb: bool,
         video_ptb: bool,
     ):
@@ -72,7 +71,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
         num_self_attns = 0
         if run_vx:
             num_self_attns += 1
-        if run_ax and audio_length > 0:
+        if run_ax:
             num_self_attns += 1
 
         video_attn_idx = 0
@@ -85,6 +84,20 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
 
         return list(stg_indexes)
 
+    def unpack_latents(self, x: torch.Tensor):
+        latent_shapes = (
+            self.conds.get("positive", {})[0]
+            .get("model_conds", {})
+            .get("latent_shapes", None)
+            .cond
+        )
+        vx, ax = comfy.utils.unpack_latents(x, latent_shapes)
+        return vx, ax
+
+    def pack_latents(self, vx: torch.Tensor, ax: torch.Tensor):
+        x, latent_shapes = comfy.utils.pack_latents([vx, ax])
+        return x, latent_shapes
+
     def predict_noise(
         self,
         x: torch.Tensor,
@@ -92,7 +105,6 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
         model_options: dict = {},
         seed=None,
     ):
-        x_shape = x.shape
         # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
         # but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
         current_step = self.current_step
@@ -108,23 +120,18 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
             Modality.VIDEO.value, GuiderParameters()
         )
 
-        audio_length = model_options["transformer_options"].get("audio_length", 0)
-
         run_vx = not video_params.do_skip(current_step)
         run_ax = not audio_params.do_skip(current_step)
         run_a2v = video_params.do_cross_attn(current_step)
         run_v2a = audio_params.do_cross_attn(current_step)
 
-        vx, ax = self.inner_model.diffusion_model.separate_audio_and_video_latents(
-            x, audio_length
-        )
+        vx, ax = self.unpack_latents(x)
         if not run_vx:
             vx = self.last_denoised_v
         if not run_ax:
             ax = self.last_denoised_a
-        x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
-            vx, ax, x_shape
-        )
+        x, _ = self.pack_latents(vx, ax)
+
         if not run_vx and not run_ax:
             return x
 
@@ -145,11 +152,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
             del model_options["transformer_options"]["run_ax"]
             del model_options["transformer_options"]["a2v_cross_attn"]
             del model_options["transformer_options"]["v2a_cross_attn"]
-        v_noise_pred_pos, a_noise_pred_pos = (
-            self.inner_model.diffusion_model.separate_audio_and_video_latents(
-                noise_pred_pos, audio_length
-            )
-        )
+        v_noise_pred_pos, a_noise_pred_pos = self.unpack_latents(noise_pred_pos)
 
         a_noise_pred_neg, v_noise_pred_neg = 0, 0
         a_noise_pred_perturbed, v_noise_pred_perturbed = 0, 0
@@ -168,11 +171,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
                     timestep,
                     model_options,
                 )[0]
-                v_noise_pred_neg, a_noise_pred_neg = (
-                    self.inner_model.diffusion_model.separate_audio_and_video_latents(
-                        noise_pred_neg, audio_length
-                    )
-                )
+                v_noise_pred_neg, a_noise_pred_neg = self.unpack_latents(noise_pred_neg)
             finally:
                 del model_options["transformer_options"]["run_vx"]
                 del model_options["transformer_options"]["run_ax"]
@@ -183,8 +182,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
             try:
                 stg_indexes = self.calc_stg_indexes(
                     run_vx,
-                    run_ax,
-                    audio_length,
+                    run_ax and ax.numel() > 0,
                     audio_params.perturb_attn,
                     video_params.perturb_attn,
                 )
@@ -202,10 +200,8 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
                     timestep,
                     model_options,
                 )[0]
-                v_noise_pred_perturbed, a_noise_pred_perturbed = (
-                    self.inner_model.diffusion_model.separate_audio_and_video_latents(
-                        noise_pred_perturbed, audio_length
-                    )
+                v_noise_pred_perturbed, a_noise_pred_perturbed = self.unpack_latents(
+                    noise_pred_perturbed
                 )
             finally:
                 self.stg_flag.do_skip = False
@@ -229,10 +225,8 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
                     timestep,
                     model_options,
                 )[0]
-                v_noise_pred_modality, a_noise_pred_modality = (
-                    self.inner_model.diffusion_model.separate_audio_and_video_latents(
-                        noise_pred_modality, audio_length
-                    )
+                v_noise_pred_modality, a_noise_pred_modality = self.unpack_latents(
+                    noise_pred_modality
                 )
             finally:
                 del model_options["transformer_options"]["a2v_cross_attn"]
@@ -260,9 +254,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
         else:
             ax = self.last_denoised_a
 
-        x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
-            vx, ax, x_shape
-        )
+        x, _ = self.pack_latents(vx, ax)
 
         # normally this would be done in cfg_function, but we skipped
         # that for efficiency: we can compute the noise predictions in
@@ -285,11 +277,7 @@ class MultimodalGuider(comfy.samplers.CFGGuider):
             }
             x = fn(args)
 
-        self.last_denoised_v, self.last_denoised_a = (
-            self.inner_model.diffusion_model.separate_audio_and_video_latents(
-                x, audio_length
-            )
-        )
+        self.last_denoised_v, self.last_denoised_a = self.unpack_latents(x)
         return x