[diffusion] feat: support LoRA for LTX2.3 (#23649)

docs_new/cookbook/diffusion/LTX/LTX.mdx

@@ -0,0 +1,206 @@
+---
+title: LTX
+description: Run LTX-2 and LTX-2.3 video generation pipelines with SGLang Diffusion.
+metatags:
+    description: "Deploy and use LTX-2 and LTX-2.3 video generation models with SGLang Diffusion, including one-stage, two-stage, HQ, TI2V, and LoRA examples."
+---
+
+import { LTXDeployment } from '/src/snippets/diffusion/ltx-deployment.jsx';
+
+## 1. Model Introduction
+
+[LTX-2](https://huggingface.co/Lightricks/LTX-2) and [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) are video generation models from Lightricks. SGLang Diffusion supports the LTX series through native one-stage and two-stage pipelines for text-to-video and image-conditioned video generation.
+
+Use `Lightricks/LTX-2` or `Lightricks/LTX-2.3` as `--model-path`. For two-stage generation, SGLang uses the spatial upsampler and distilled LoRA components from the model snapshot by default. LTX-2.3 also supports the HQ two-stage variant.
+
+<Warning>
+**License notice:** LTX-2 and LTX-2.3 are released under the LTX-2 Community License Agreement, not Apache 2.0. The license includes commercial-use restrictions for some entities. Review the [official Lightricks license](https://huggingface.co/Lightricks/LTX-2.3/blob/main/LICENSE) before production or commercial use; SGLang support does not grant additional model usage rights.
+</Warning>
+
+## 2. SGLang-diffusion Installation
+
+Install SGLang with diffusion dependencies:
+
+```bash Command
+uv pip install "sglang[diffusion]" --prerelease=allow
+```
+
+For platform-specific setup, see the [SGLang Diffusion installation guide](/docs/sglang-diffusion/installation).
+
+## 3. Model Deployment
+
+This section provides deployment configurations optimized for different LTX pipelines and hardware targets.
+
+### 3.1 Basic Configuration
+
+The LTX series supports one-stage and two-stage pipelines. LTX-2.3 also supports the HQ two-stage pipeline. The recommended launch configuration depends on whether the target GPU can keep both two-stage DiTs resident.
+
+**Interactive Command Generator**: Use the configuration selector below to generate a deployment command. The default selection targets a single NVIDIA H200 with `resident` two-stage mode, which is the fastest startup path for the specified high-memory environment.
+
+<LTXDeployment />
+
+### 3.2 Configuration Tips
+
+Choose the pipeline class based on the quality and latency target:
+
+| Use case | Pipeline class | Notes |
+| --- | --- | --- |
+| One-stage generation | `LTX2Pipeline` | Fastest LTX native path. Supports T2V and TI2V. |
+| Two-stage generation | `LTX2TwoStagePipeline` | Uses a base stage and a refinement stage. Supported by LTX-2 and LTX-2.3. |
+| Two-stage High Quality (HQ) generation | `LTX2TwoStageHQPipeline` | LTX-2.3 HQ path; defaults to 1920x1088 unless you override `--width` and `--height`. |
+
+Feature compatibility:
+
+| Pipeline class | T2V | TI2V (`--image-path`) | LoRA (`--lora-path`) | Notes |
+| --- | --- | --- | --- | --- |
+| `LTX2Pipeline` | Yes | Yes | Yes | One-stage path. Cannot be combined with HQ because HQ is a separate two-stage pipeline class. |
+| `LTX2TwoStagePipeline` | Yes | Yes | Yes | Standard two-stage path for LTX-2 and LTX-2.3. |
+| `LTX2TwoStageHQPipeline` | Yes | Yes | Yes | High Quality two-stage path for LTX-2.3. Use this instead of `LTX2Pipeline`; it is not a one-stage mode flag. |
+
+For two-stage pipelines, `--ltx2-two-stage-device-mode` controls transformer residency:
+
+| Mode | When to use it |
+| --- | --- |
+| `snapshot` | Recommended default. Balances latency and VRAM. |
+| `resident` | Best latency on high-VRAM GPUs because both DiTs can stay resident. |
+| `original` | Closest to the original two-stage switching semantics. |
+
+Other deployment flags:
+
+- `--lora-path`: Preload a community LoRA adapter.
+- `--lora-weight-name`: Select the exact safetensors file when the LoRA repository contains multiple weight files.
+
+<Note>
+For native LTX-2.3 two-stage serving without a user LoRA, `resident` is the fastest high-VRAM path. When you pass `--lora-path`, SGLang still applies the user LoRA during the two-stage switch, so use `resident` on H200-class GPUs for enough VRAM, but do not expect the same premerged-stage2 benefit as the no-user-LoRA path.
+</Note>
+
+## 4. Model Invocation
+
+### 4.1 Basic Usage
+
+The examples below spell out the current SGLang sampling defaults for reproducibility:
+
+| Model path | Default output | Default frames | Default steps |
+| --- | --- | --- | --- |
+| `Lightricks/LTX-2` | 768x512 | 121 | 40 |
+| `Lightricks/LTX-2.3` | 768x512 | 121 | 30 |
+| `Lightricks/LTX-2.3` with `LTX2TwoStageHQPipeline` | 1920x1088 | 121 | 15 |
+
+#### 4.1.1 LTX-2 one-stage text-to-video
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2 \
+  --pipeline-class-name LTX2Pipeline \
+  --prompt "A quiet coastal town at sunrise, fishing boats moving slowly through golden mist, cinematic camera movement" \
+  --save-output
+```
+
+#### 4.1.2 LTX-2.3 one-stage text-to-video
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2Pipeline \
+  --prompt "A quiet coastal town at sunrise, fishing boats moving slowly through golden mist, cinematic camera movement" \
+  --save-output
+```
+
+#### 4.1.3 LTX-2 two-stage text-to-video
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --prompt "A handheld shot follows a red tram crossing a rainy city square at night, reflections on the pavement, cinematic lighting" \
+  --save-output
+```
+
+#### 4.1.4 LTX-2.3 two-stage text-to-video
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --prompt "A handheld shot follows a red tram crossing a rainy city square at night, reflections on the pavement, cinematic lighting" \
+  --save-output
+```
+
+#### 4.1.5 LTX-2.3 HQ text-to-video
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStageHQPipeline \
+  --prompt "A wide cinematic shot of alpine clouds rolling over a mountain ridge, soft morning light, slow aerial camera movement" \
+  --save-output
+```
+
+#### 4.1.6 Image-to-video with one reference image
+
+Pass one image to `--image-path` for image-conditioned generation:
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --image-path ./inputs/start.png \
+  --prompt "The camera slowly pushes forward as the subject turns toward warm window light, subtle natural motion, cinematic" \
+  --save-output
+```
+
+#### 4.1.7 First-to-last-frame transition with two reference images
+
+Pass two images to `--image-path` for transition-style TI2V. The first image is used as the starting condition and the second image is used as the ending condition.
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --image-path ./inputs/start.png ./inputs/end.png \
+  --prompt "A smooth cinematic transition from the first scene into the final scene, dynamic camera motion, motion blur, zhuanchang" \
+  --save-output
+```
+
+### 4.2 Advanced Usage
+
+#### 4.2.1 Use community LoRAs
+
+Use `--lora-path` to load a LoRA adapter. If the Hugging Face repo contains multiple safetensors files, use `--lora-weight-name` to select the exact file. `--lora-scale` maps to the standard LoRA merge scale and defaults to `1.0`.
+
+The following example uses [`valiantcat/LTX-2.3-Transition-LORA`](https://huggingface.co/valiantcat/LTX-2.3-Transition-LORA):
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --lora-path valiantcat/LTX-2.3-Transition-LORA \
+  --lora-weight-name ltx2.3-transition.safetensors \
+  --prompt "A low-angle tracking shot moves through a foggy forest road. The camera rises above the treetops and transitions into a clear view of a snowy mountain peak under bright sunlight, zhuanchang" \
+  --save-output
+```
+
+You can combine the Transition LoRA with two reference images:
+
+```bash Command
+sglang generate \
+  --model-path Lightricks/LTX-2.3 \
+  --pipeline-class-name LTX2TwoStagePipeline \
+  --image-path ./inputs/start.png ./inputs/end.png \
+  --lora-path valiantcat/LTX-2.3-Transition-LORA \
+  --lora-weight-name ltx2.3-transition.safetensors \
+  --prompt "A fast cinematic transition from the first image to the second image, whip-pan motion, atmospheric lighting, zhuanchang" \
+  --save-output
+```
+
+<Note>
+Some community LoRAs only include weights for transformer blocks. In that case, SGLang logs a concise coverage summary and leaves unmatched LoRA-capable layers on the base model weights. This is expected when the adapter format intentionally omits those layers.
+</Note>
+
+## 5. Practical Tips
+
+- Use `--pipeline-class-name LTX2TwoStagePipeline` as the default LTX two-stage quality path.
+- Use `--pipeline-class-name LTX2TwoStageHQPipeline` when you want the HQ path and have enough VRAM for larger outputs.
+- Use `--ltx2-two-stage-device-mode resident` on high-VRAM GPUs if latency matters more than memory usage.
+- Use `--ltx2-two-stage-device-mode original` when comparing against official two-stage behavior.
+- Keep `--width` and `--height` aligned with the target model resolution; for LTX models, these are output video dimensions.

docs_new/cookbook/diffusion/intro.mdx

@@ -19,6 +19,12 @@ metatags:
     href="/cookbook/diffusion/Wan/Wan2.2"
     img="/cards/logos/wan.png"
   />
+  <Card
+    title="LTX"
+    mode="card"
+    href="/cookbook/diffusion/LTX/LTX"
+    img="/cards/Diffusion-card.png"
+  />
   <Card
     title="Qwen-Image"
     mode="card"

docs_new/docs.json

@@ -1094,6 +1094,12 @@
                       "cookbook/diffusion/Wan/Wan2.2"
                     ]
                   },
+                  {
+                    "group": "LTX",
+                    "pages": [
+                      "cookbook/diffusion/LTX/LTX"
+                    ]
+                  },
                   {
                     "group": "Qwen-Image",
                     "pages": [

docs_new/src/snippets/diffusion/ltx-deployment.jsx

@@ -0,0 +1,233 @@
+export const LTXDeployment = () => {
+  const options = {
+    hardware: {
+      name: 'hardware',
+      title: 'Hardware Platform',
+      items: [
+        { id: 'h200', label: 'H200', subtitle: 'Fastest, resident', default: true },
+        { id: 'standard', label: 'Standard CUDA', subtitle: 'Snapshot mode', default: false },
+        { id: 'official', label: 'Official Match', subtitle: 'Original switching', default: false },
+      ],
+    },
+    model: {
+      name: 'model',
+      title: 'Model',
+      items: [
+        { id: 'ltx23', label: 'LTX-2.3', default: true },
+        { id: 'ltx2', label: 'LTX-2', default: false },
+      ],
+    },
+    pipeline: {
+      name: 'pipeline',
+      title: 'Pipeline',
+      items: [
+        { id: 'two-stage', label: 'Two Stage', default: true, validModels: ['ltx2', 'ltx23'] },
+        { id: 'two-stage-hq', label: 'Two Stage HQ', subtitle: 'High Quality', default: false, validModels: ['ltx23'] },
+        { id: 'one-stage', label: 'One Stage', default: false, validModels: ['ltx2', 'ltx23'] },
+      ],
+    },
+  };
+
+  const modelConfigs = {
+    ltx2: {
+      repoId: 'Lightricks/LTX-2',
+      pipelines: {
+        'one-stage': 'LTX2Pipeline',
+        'two-stage': 'LTX2TwoStagePipeline',
+      },
+      supportedLoras: [],
+    },
+    ltx23: {
+      repoId: 'Lightricks/LTX-2.3',
+      pipelines: {
+        'one-stage': 'LTX2Pipeline',
+        'two-stage': 'LTX2TwoStagePipeline',
+        'two-stage-hq': 'LTX2TwoStageHQPipeline',
+      },
+      supportedLoras: [
+        {
+          id: 'transition',
+          path: 'valiantcat/LTX-2.3-Transition-LORA',
+          weightName: 'ltx2.3-transition.safetensors',
+          validPipelines: ['two-stage', 'two-stage-hq'],
+        },
+      ],
+    },
+  };
+
+  const getInitialState = () => ({
+    hardware: 'h200',
+    model: 'ltx23',
+    pipeline: 'two-stage',
+    selectedLoraPath: 'none',
+  });
+
+  const [values, setValues] = useState(getInitialState);
+  const [isDark, setIsDark] = useState(false);
+
+  useEffect(() => {
+    const checkDarkMode = () => {
+      const html = document.documentElement;
+      const isDarkMode = html.classList.contains('dark') ||
+        html.getAttribute('data-theme') === 'dark' ||
+        html.style.colorScheme === 'dark';
+      setIsDark(isDarkMode);
+    };
+    checkDarkMode();
+    const observer = new MutationObserver(checkDarkMode);
+    observer.observe(document.documentElement, { attributes: true, attributeFilter: ['class', 'data-theme', 'style'] });
+    return () => observer.disconnect();
+  }, []);
+
+  const availableLoras = (() => {
+    const config = modelConfigs[values.model];
+    return (config?.supportedLoras || []).filter((lora) => lora.validPipelines.includes(values.pipeline));
+  })();
+
+  const handleRadioChange = (optionName, itemId) => {
+    setValues((prev) => {
+      const next = { ...prev, [optionName]: itemId };
+
+      const validPipeline = options.pipeline.items.some((item) => (
+        item.id === next.pipeline && item.validModels.includes(next.model)
+      ));
+      if (!validPipeline) {
+        next.pipeline = 'two-stage';
+      }
+
+      const config = modelConfigs[next.model];
+      const nextSupported = (config?.supportedLoras || []).filter((lora) => lora.validPipelines.includes(next.pipeline));
+      const isValid = nextSupported.some((lora) => lora.path === prev.selectedLoraPath);
+      if (!isValid) {
+        next.selectedLoraPath = 'none';
+      }
+      return next;
+    });
+  };
+
+  const handleLoraToggle = (path) => {
+    setValues((prev) => ({
+      ...prev,
+      selectedLoraPath: prev.selectedLoraPath === path ? 'none' : path,
+    }));
+  };
+
+  const getDeviceMode = () => {
+    if (values.hardware === 'h200') {
+      return 'resident';
+    }
+    if (values.hardware === 'official') {
+      return 'original';
+    }
+    return 'snapshot';
+  };
+
+  const generateCommand = () => {
+    const config = modelConfigs[values.model];
+    const pipelineClass = config.pipelines[values.pipeline];
+    if (!pipelineClass) {
+      return '# Error: Invalid configuration';
+    }
+
+    let command = `sglang serve \\\n  --model-path ${config.repoId} \\\n  --pipeline-class-name ${pipelineClass}`;
+    if (values.pipeline !== 'one-stage') {
+      command += ` \\\n  --ltx2-two-stage-device-mode ${getDeviceMode()}`;
+    }
+
+    const selectedLora = availableLoras.find((lora) => lora.path === values.selectedLoraPath);
+    if (selectedLora) {
+      command += ` \\\n  --lora-path ${selectedLora.path} \\\n  --lora-weight-name ${selectedLora.weightName}`;
+    }
+
+    command += ` \\\n  --port 30000`;
+    return command;
+  };
+
+  const containerStyle = { maxWidth: '900px', margin: '0 auto', display: 'flex', flexDirection: 'column', gap: '4px' };
+  const cardStyle = { padding: '8px 12px', border: `1px solid ${isDark ? '#374151' : '#e5e7eb'}`, borderLeft: `3px solid ${isDark ? '#E85D4D' : '#D45D44'}`, borderRadius: '4px', display: 'flex', alignItems: 'center', gap: '12px', background: isDark ? '#1f2937' : '#fff' };
+  const titleStyle = { fontSize: '13px', fontWeight: '600', minWidth: '140px', flexShrink: 0, color: isDark ? '#e5e7eb' : 'inherit' };
+  const itemsStyle = { display: 'flex', rowGap: '2px', columnGap: '6px', flexWrap: 'wrap', alignItems: 'center', flex: 1 };
+  const labelBaseStyle = { padding: '4px 10px', border: `1px solid ${isDark ? '#9ca3af' : '#d1d5db'}`, borderRadius: '3px', cursor: 'pointer', display: 'inline-flex', flexDirection: 'column', alignItems: 'center', justifyContent: 'center', fontWeight: '500', fontSize: '13px', transition: 'all 0.2s', userSelect: 'none', minWidth: '45px', textAlign: 'center', flex: 1, background: isDark ? '#374151' : '#fff', color: isDark ? '#e5e7eb' : 'inherit' };
+  const checkedStyle = { background: '#D45D44', color: 'white', borderColor: '#D45D44' };
+  const subtitleStyle = { display: 'block', fontSize: '9px', marginTop: '1px', lineHeight: '1.1', opacity: 0.7 };
+  const commandDisplayStyle = { flex: 1, padding: '12px 16px', background: isDark ? '#111827' : '#f5f5f5', borderRadius: '6px', fontFamily: "'Menlo', 'Monaco', 'Courier New', monospace", fontSize: '12px', lineHeight: '1.5', color: isDark ? '#e5e7eb' : '#374151', whiteSpace: 'pre-wrap', overflowX: 'auto', margin: 0, border: `1px solid ${isDark ? '#374151' : '#e5e7eb'}` };
+
+  return (
+    <div style={containerStyle} className="not-prose">
+      {Object.entries(options).map(([key, option]) => {
+        const itemsToDisplay = key === 'pipeline'
+          ? option.items.filter((item) => item.validModels.includes(values.model))
+          : option.items;
+
+        return (
+          <div key={key} style={cardStyle}>
+            <div style={titleStyle}>{option.title}</div>
+            <div style={itemsStyle}>
+              {itemsToDisplay.map((item) => {
+                const isChecked = values[option.name] === item.id;
+                return (
+                  <label key={item.id} style={{ ...labelBaseStyle, ...(isChecked ? checkedStyle : {}) }}>
+                    <input
+                      type="radio"
+                      name={option.name}
+                      checked={isChecked}
+                      onChange={() => handleRadioChange(key, item.id)}
+                      style={{ display: 'none' }}
+                    />
+                    {item.label}
+                    {item.subtitle && (
+                      <small style={{ ...subtitleStyle, color: isChecked ? 'rgba(255,255,255,0.85)' : 'inherit' }}>
+                        {item.subtitle}
+                      </small>
+                    )}
+                  </label>
+                );
+              })}
+            </div>
+          </div>
+        );
+      })}
+
+      <div style={cardStyle}>
+        <div style={titleStyle}>Select LoRA Model</div>
+        <div style={itemsStyle}>
+          {availableLoras.length === 0 && (
+            <div style={{ color: isDark ? '#999' : '#666', fontSize: '12px', padding: '8px' }}>
+              No LoRA models available for this configuration.
+            </div>
+          )}
+          {availableLoras.map((lora) => {
+            const isSelected = values.selectedLoraPath === lora.path;
+            return (
+              <label
+                key={lora.id}
+                style={{ ...labelBaseStyle, ...(isSelected ? checkedStyle : {}) }}
+                onClick={(event) => {
+                  event.preventDefault();
+                  handleLoraToggle(lora.path);
+                }}
+              >
+                <input
+                  type="radio"
+                  name="loraModelSelection"
+                  checked={isSelected}
+                  readOnly
+                  style={{ display: 'none' }}
+                />
+                {lora.id}
+                <small style={{ ...subtitleStyle, color: isSelected ? 'rgba(255,255,255,0.85)' : 'inherit' }}>
+                  {lora.path}
+                </small>
+              </label>
+            );
+          })}
+        </div>
+      </div>
+
+      <div style={cardStyle}>
+        <div style={titleStyle}>Run this Command:</div>
+        <pre style={commandDisplayStyle}>{generateCommand()}</pre>
+      </div>
+    </div>
+  );
+};

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -2,7 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
-
+import os
 
 import torch
 from torch import nn
@@ -37,10 +37,17 @@ torch._dynamo.config.recompile_limit = 64
 
 
 LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024
+LoRAWeightEntry = tuple[
+    torch.nn.Parameter,
+    torch.nn.Parameter,
+    str | None,
+    float,
+    int | None,
+    int | None,
+]
 
 
 class BaseLayerWithLoRA(nn.Module):
-
     def __init__(
         self,
         base_layer: nn.Module,
@@ -60,9 +67,7 @@ class BaseLayerWithLoRA(nn.Module):
         self.disable_lora: bool = True
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
-        self.lora_weights_list: list[
-            tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
-        ] = []
+        self.lora_weights_list: list[LoRAWeightEntry] = []
         self.lora_path: str | None = None
         self.strength: float = 1.0
 
@@ -147,7 +152,16 @@ class BaseLayerWithLoRA(nn.Module):
             self.strength = 1.0
 
         # Add to list for multi-LoRA support
-        self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength))
+        self.lora_weights_list.append(
+            (
+                lora_A_param,
+                lora_B_param,
+                lora_path,
+                strength,
+                self.lora_rank,
+                self.lora_alpha,
+            )
+        )
 
         # Set backward compatibility attributes to point to the last LoRA (for single LoRA case)
         # This ensures backward compatibility while supporting multiple LoRA
@@ -166,29 +180,27 @@ class BaseLayerWithLoRA(nn.Module):
     def _merge_lora_into_data(
         self,
         data: torch.Tensor,
-        lora_list: list[
-            tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
-        ],
+        lora_list: list[LoRAWeightEntry],
     ) -> None:
         """
         Merge all LoRA adapters into the data tensor in-place.
 
         Args:
             data: The base weight tensor to merge LoRA into (modified in-place)
-            lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples
+            lora_list: List of (lora_A, lora_B, lora_path, lora_strength, rank, alpha) tuples
         """
         # Merge all LoRA adapters in order
-        for lora_A, lora_B, _, lora_strength in lora_list:
+        for lora_A, lora_B, _, lora_strength, lora_rank, lora_alpha in lora_list:
             lora_A_sliced = self.slice_lora_a_weights(lora_A.to(data))
             lora_B_sliced = self.slice_lora_b_weights(lora_B.to(data))
 
             scale = lora_strength
             if (
-                self.lora_alpha is not None
-                and self.lora_rank is not None
-                and self.lora_alpha != self.lora_rank
+                lora_alpha is not None
+                and lora_rank is not None
+                and lora_alpha != lora_rank
             ):
-                scale *= self.lora_alpha / self.lora_rank
+                scale *= lora_alpha / lora_rank
 
             if not isinstance(lora_B_sliced, torch.Tensor):
                 lora_delta = lora_B_sliced @ lora_A_sliced
@@ -222,6 +234,17 @@ class BaseLayerWithLoRA(nn.Module):
                 chunk_delta = lora_B_2d[start:end] @ lora_A_sliced
                 data_2d[start:end].add_(chunk_delta, alpha=scale)
 
+    def _should_merge_in_fp32(
+        self,
+        lora_list: list[LoRAWeightEntry],
+    ) -> bool:
+        if os.getenv("SGLANG_DIFFUSION_LORA_MERGE_FP32", "0") != "1":
+            return False
+        for _, _, lora_path, _, _, _ in lora_list:
+            if lora_path and "distilled-lora" in lora_path.lower():
+                return False
+        return True
+
     @torch.no_grad()
     def merge_lora_weights(self, strength: float | None = None) -> None:
         if strength is not None:
@@ -236,11 +259,22 @@ class BaseLayerWithLoRA(nn.Module):
         # Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility
         lora_list = self.lora_weights_list if self.lora_weights_list else []
         if not lora_list and self.lora_A is not None and self.lora_B is not None:
-            lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)]
+            lora_list = [
+                (
+                    self.lora_A,
+                    self.lora_B,
+                    self.lora_path,
+                    self.strength,
+                    self.lora_rank,
+                    self.lora_alpha,
+                )
+            ]
 
         if not lora_list:
             raise ValueError("LoRA weights not set. Please set them first.")
 
+        merge_in_fp32 = self._should_merge_in_fp32(lora_list)
+
         if isinstance(self.base_layer.weight, DTensor):
             mesh = self.base_layer.weight.data.device_mesh
             unsharded_base_layer = ReplicatedLinear(
@@ -257,10 +291,19 @@ class BaseLayerWithLoRA(nn.Module):
             data = self.base_layer.weight.data.to(
                 get_local_torch_device()
             ).full_tensor()
+            target_dtype = data.dtype
+            if (
+                merge_in_fp32
+                and data.is_floating_point()
+                and data.dtype != torch.float32
+            ):
+                data = data.to(torch.float32)
 
             self._merge_lora_into_data(data, lora_list)
 
-            unsharded_base_layer.weight = nn.Parameter(data.to(current_device))
+            unsharded_base_layer.weight = nn.Parameter(
+                data.to(current_device, dtype=target_dtype)
+            )
             if isinstance(getattr(self.base_layer, "bias", None), DTensor):
                 unsharded_base_layer.bias = nn.Parameter(
                     self.base_layer.bias.to(get_local_torch_device(), non_blocking=True)
@@ -282,10 +325,19 @@ class BaseLayerWithLoRA(nn.Module):
         else:
             current_device = self.base_layer.weight.data.device
             data = self.base_layer.weight.data.to(get_local_torch_device())
+            target_dtype = data.dtype
+            if (
+                merge_in_fp32
+                and data.is_floating_point()
+                and data.dtype != torch.float32
+            ):
+                data = data.to(torch.float32)
 
             self._merge_lora_into_data(data, lora_list)
 
-            self.base_layer.weight.data = data.to(current_device, non_blocking=True)
+            self.base_layer.weight.data = data.to(
+                current_device, dtype=target_dtype, non_blocking=True
+            )
 
         self.merged = True
 
@@ -342,7 +394,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
 
 
 class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
-
     def __init__(
         self,
         base_layer: ColumnParallelLinear,
@@ -400,7 +451,6 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
 
 
 class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
-
     def __init__(
         self,
         base_layer: MergedColumnParallelLinear,
@@ -422,7 +472,6 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
 
 
 class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
-
     def __init__(
         self,
         base_layer: QKVParallelLinear,
@@ -455,7 +504,6 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
 
 
 class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
-
     def __init__(
         self,
         base_layer: RowParallelLinear,

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -494,19 +494,25 @@ class LTX2TwoStageDeviceManager:
         if "stage2" in self._phase_ready_events:
             return
         if self._snapshot_low_vram_mode:
-            stage1_module = self.pipeline.get_module("transformer")
-            stage1_param = (
-                next(stage1_module.parameters(), None)
-                if stage1_module is not None
-                else None
-            )
-            if stage1_param is not None and stage1_param.device.type == "cuda":
-                self._release_module_to_cpu_snapshot("transformer")
+            self._release_stage1_for_low_vram()
 
         self._schedule_phase_prefetch(
             "stage2", self.pipeline.get_module("transformer_2")
         )
 
+    def prepare_upsample_after_stage1(self) -> bool:
+        if (
+            not self.should_use_premerged
+            or self.mode != "snapshot"
+            or not self.server_args.dit_cpu_offload
+            or not self._snapshot_low_vram_mode
+        ):
+            return False
+        if "stage2" in self._phase_ready_events:
+            return False
+        self._release_stage1_for_low_vram()
+        return True
+
     def ensure_phase_ready(self, phase: str | None) -> None:
         if not self.should_use_premerged or phase not in ("stage1", "stage2"):
             return
@@ -616,6 +622,16 @@ class LTX2TwoStageDeviceManager:
         phase = "stage2" if module_name == "transformer_2" else "stage1"
         self._phase_ready_events.pop(phase, None)
 
+    def _release_stage1_for_low_vram(self) -> None:
+        stage1_module = self.pipeline.get_module("transformer")
+        stage1_param = (
+            next(stage1_module.parameters(), None)
+            if stage1_module is not None
+            else None
+        )
+        if stage1_param is not None and stage1_param.device.type == "cuda":
+            self._release_module_to_cpu_snapshot("transformer")
+
     def _ensure_on_gpu(self, module_name: str) -> None:
         module = self.pipeline.get_module(module_name)
         if module is None:
@@ -807,6 +823,9 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
     def prefetch_ltx2_stage2_after_stage1(self) -> None:
         self._device_manager.prefetch_stage2_after_stage1()
 
+    def prepare_ltx2_upsample_after_stage1(self) -> bool:
+        return self._device_manager.prepare_upsample_after_stage1()
+
     def should_skip_ltx2_lora_switch_stage(self) -> bool:
         return self._use_premerged_stage2_transformer and self._device_manager.mode in (
             "snapshot",

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -99,7 +99,9 @@ class LoRAPipeline(ComposedPipelineBase):
         if self.lora_path is not None:
             self.convert_to_lora_layers()
             self.set_lora(
-                self.lora_nickname, self.lora_path, strength=self.server_args.lora_scale  # type: ignore
+                self.lora_nickname,
+                self.lora_path,
+                strength=self.server_args.lora_scale,  # type: ignore
             )  # type: ignore
 
     def is_target_layer(self, module_name: str) -> bool:
@@ -426,6 +428,8 @@ class LoRAPipeline(ComposedPipelineBase):
             )
 
         adapted_count = 0
+        missing_layers_by_adapter = [[] for _ in lora_nicknames]
+        applied_count_by_adapter = [0 for _ in lora_nicknames]
         for name, layer in lora_layers.items():
             # Apply all LoRA adapters in order
             for idx, (nickname, path, lora_strength) in enumerate(
@@ -465,13 +469,9 @@ class LoRAPipeline(ComposedPipelineBase):
                         ),  # Only clear on first LoRA
                     )
                     adapted_count += 1
+                    applied_count_by_adapter[idx] += 1
                 else:
-                    if rank == 0 and idx == 0:  # Only warn for first missing LoRA
-                        logger.warning(
-                            "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.",
-                            path,
-                            name,
-                        )
+                    missing_layers_by_adapter[idx].append(name)
                     # Only disable if no LoRA was applied at all
                     if idx == len(lora_nicknames) - 1:
                         has_any_lora = any(
@@ -481,6 +481,37 @@ class LoRAPipeline(ComposedPipelineBase):
                         )
                         if not has_any_lora:
                             layer.disable_lora = True
+
+        if rank == 0:
+            total_layers = len(lora_layers)
+            example_limit = 8
+            for idx, path in enumerate(lora_paths):
+                missing_layers = missing_layers_by_adapter[idx]
+                if not missing_layers:
+                    continue
+                missing_count = len(missing_layers)
+                applied_count = applied_count_by_adapter[idx]
+                examples = ", ".join(missing_layers[:example_limit])
+                if missing_count > example_limit:
+                    examples += ", ..."
+                if applied_count == 0:
+                    logger.warning(
+                        "LoRA adapter %s did not match any LoRA layer. "
+                        "Checked %d layers; examples: %s",
+                        path,
+                        total_layers,
+                        examples,
+                    )
+                else:
+                    logger.info(
+                        "LoRA adapter %s covers %d/%d LoRA layers; "
+                        "%d layers use base weights. Examples: %s",
+                        path,
+                        applied_count,
+                        total_layers,
+                        missing_count,
+                        examples,
+                    )
         return adapted_count
 
     def is_lora_effective(self, target: str = "all") -> bool:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -217,6 +217,14 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Run the distilled refinement schedule on top of the shared AV denoiser."""
         batch.extra["ltx2_phase"] = "stage2"
+        pipeline = self.pipeline() if self.pipeline else None
+        ensure_phase_ready = (
+            getattr(pipeline, "ensure_ltx2_phase_ready", None)
+            if pipeline is not None
+            else None
+        )
+        if callable(ensure_phase_ready):
+            ensure_phase_ready("stage2")
         original_clean_latent_background = getattr(
             batch, "ltx2_ti2v_clean_latent_background", None
         )
@@ -250,6 +258,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
             renoise_generator = None
         if is_native_ti2v:
             prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state(
+                batch=batch,
                 latents=batch.latents,
                 image_latent=batch.image_latent,
                 num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -417,25 +417,48 @@ class LTX2ImageEncodingStage(PipelineStage):
         ):
             return self._condition_image_encoder(video_condition)
 
+    @staticmethod
+    def _normalize_ltx2_image_paths(image_path: str | list[str]) -> list[str]:
+        image_paths = image_path if isinstance(image_path, list) else [image_path]
+        if len(image_paths) > 2:
+            raise ValueError(
+                "LTX-2 TI2V currently supports at most two conditioning images "
+                "([first_frame, last_frame])."
+            )
+        return image_paths
+
+    @staticmethod
+    def _normalize_ltx2_image_latents(
+        image_latent: torch.Tensor | list[torch.Tensor] | None,
+    ) -> list[torch.Tensor]:
+        if image_latent is None:
+            return []
+        return image_latent if isinstance(image_latent, list) else [image_latent]
+
     # -- forward ---------------------------------------------------------
 
     @torch.no_grad()
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         if batch.image_path is None:
             return batch
+        image_paths = self._normalize_ltx2_image_paths(batch.image_path)
+
+        vae_sf = int(server_args.pipeline_config.vae_scale_factor)
+        patch = int(server_args.pipeline_config.patch_size)
+        expected_tokens = (int(batch.height) // vae_sf // patch) * (
+            int(batch.width) // vae_sf // patch
+        )
         if (
             batch.image_latent is not None
             and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0
         ):
             # Re-encode if resolution changed (e.g. two-stage upsample between stages)
-            vae_sf = int(server_args.pipeline_config.vae_scale_factor)
-            patch = int(server_args.pipeline_config.patch_size)
-            expected = (int(batch.height) // vae_sf // patch) * (
-                int(batch.width) // vae_sf // patch
-            )
-            if int(batch.image_latent.shape[1]) == expected:
+            existing_latents = self._normalize_ltx2_image_latents(batch.image_latent)
+            if len(existing_latents) == len(image_paths) and all(
+                int(latent.shape[1]) == expected_tokens for latent in existing_latents
+            ):
                 return batch
-            # Resolution mismatch — clear and re-encode below
+            # Resolution or reference-count mismatch — clear and re-encode below
             batch.image_latent = None
             batch.ltx2_num_image_tokens = 0
 
@@ -447,20 +470,23 @@ class LTX2ImageEncodingStage(PipelineStage):
 
         from sglang.multimodal_gen.runtime.models.vision_utils import load_image
 
-        # 1. Load image, apply codec compression, resize for condition_image
-        image_path = (
-            batch.image_path[0]
-            if isinstance(batch.image_path, list)
-            else batch.image_path
-        )
-        img = load_image(image_path)
-        arr = np.array(img).astype(np.uint8)[..., :3]
-        arr = self._apply_video_codec_compression(arr, crf=33)
-        conditioned_img = PIL.Image.fromarray(arr)
-        batch.condition_image = conditioned_img.resize(
-            (int(batch.width), int(batch.height)),
-            resample=PIL.Image.Resampling.BILINEAR,
-        )
+        # 1. Load images, apply codec compression, resize for condition_image
+        conditioned_imgs = []
+        for image_path in image_paths:
+            img = load_image(image_path)
+            arr = np.array(img).astype(np.uint8)[..., :3]
+            arr = self._apply_video_codec_compression(arr, crf=33)
+            conditioned_img = PIL.Image.fromarray(arr)
+            conditioned_imgs.append(conditioned_img)
+        batch.condition_image = [
+            img.resize(
+                (int(batch.width), int(batch.height)),
+                resample=PIL.Image.Resampling.BILINEAR,
+            )
+            for img in conditioned_imgs
+        ]
+        if len(batch.condition_image) == 1:
+            batch.condition_image = batch.condition_image[0]
 
         # 2. Load encoder(s) to device, cast to encode_dtype
         use_condition_encoder = self._ensure_condition_image_encoder(server_args)
@@ -478,21 +504,35 @@ class LTX2ImageEncodingStage(PipelineStage):
         else:
             self.vae = self.vae.to(dtype=encode_dtype)
 
-        video_condition = self._pil_to_video_tensor(
-            conditioned_img,
-            width=int(batch.width),
-            height=int(batch.height),
-            device=device,
-            dtype=encode_dtype,
-        )
-
-        # 3. Encode
-        if use_condition_encoder:
-            latent = self._condition_encode(video_condition, server_args).to(
-                dtype=encode_dtype
+        packed_latents = []
+        for conditioned_img in conditioned_imgs:
+            video_condition = self._pil_to_video_tensor(
+                conditioned_img,
+                width=int(batch.width),
+                height=int(batch.height),
+                device=device,
+                dtype=encode_dtype,
             )
-        else:
-            latent = self._vae_encode(video_condition, server_args, batch.generator)
+
+            # 3. Encode
+            if use_condition_encoder:
+                latent = self._condition_encode(video_condition, server_args).to(
+                    dtype=encode_dtype
+                )
+            else:
+                latent = self._vae_encode(video_condition, server_args, batch.generator)
+
+            packed = server_args.pipeline_config.maybe_pack_latents(
+                latent, latent.shape[0], batch
+            )
+            if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
+                raise ValueError("Expected packed image latents [B, S0, D].")
+            if int(packed.shape[1]) != expected_tokens:
+                raise ValueError(
+                    f"LTX-2 conditioning token count mismatch: "
+                    f"{packed.shape[1]=} {expected_tokens=}."
+                )
+            packed_latents.append(packed)
 
         # Restore VAE to its config dtype (shared with decoding stage)
         if not use_condition_encoder:
@@ -501,32 +541,16 @@ class LTX2ImageEncodingStage(PipelineStage):
             ]
             self.vae = self.vae.to(dtype=original_dtype)
 
-        # 4. Pack into token latents and validate
-        packed = server_args.pipeline_config.maybe_pack_latents(
-            latent, latent.shape[0], batch
+        batch.image_latent = (
+            packed_latents[0] if len(packed_latents) == 1 else packed_latents
         )
-        if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
-            raise ValueError("Expected packed image latents [B, S0, D].")
-
-        vae_sf = int(server_args.pipeline_config.vae_scale_factor)
-        patch = int(server_args.pipeline_config.patch_size)
-        expected_tokens = (int(batch.height) // vae_sf // patch) * (
-            int(batch.width) // vae_sf // patch
-        )
-        if int(packed.shape[1]) != expected_tokens:
-            raise ValueError(
-                f"LTX-2 conditioning token count mismatch: "
-                f"{packed.shape[1]=} {expected_tokens=}."
-            )
-
-        batch.image_latent = packed
-        batch.ltx2_num_image_tokens = int(packed.shape[1])
+        batch.ltx2_num_image_tokens = int(packed_latents[0].shape[1])
 
         if batch.debug:
             logger.info(
                 "LTX2 TI2V: %d tokens (shape=%s) for %sx%s",
                 batch.ltx2_num_image_tokens,
-                tuple(batch.image_latent.shape),
+                tuple(packed_latents[0].shape),
                 batch.width,
                 batch.height,
             )
@@ -703,7 +727,6 @@ class ImageVAEEncodingStage(PipelineStage):
         self,
         image: torch.Tensor | PIL.Image.Image,
     ) -> torch.Tensor:
-
         if isinstance(image, PIL.Image.Image):
             image = pil_to_numpy(image)  # to np
             image = numpy_to_pt(image)  # to pt

python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py

@@ -516,24 +516,111 @@ class LTX2DenoisingStage(DenoisingStage):
         return next_video, next_audio
 
     @staticmethod
-    def _prepare_ltx2_ti2v_clean_state(
+    def _normalize_ltx2_condition_latents(
+        image_latent: torch.Tensor | list[torch.Tensor] | None,
+    ) -> list[torch.Tensor]:
+        if image_latent is None:
+            return []
+        return image_latent if isinstance(image_latent, list) else [image_latent]
+
+    @classmethod
+    def _get_ltx2_condition_spans(
+        cls,
+        batch: Req,
         latents: torch.Tensor,
-        image_latent: torch.Tensor,
+        image_latent: torch.Tensor | list[torch.Tensor] | None,
+        num_img_tokens: int,
+    ) -> list[tuple[int, torch.Tensor]]:
+        if num_img_tokens <= 0:
+            return []
+        if not (isinstance(latents, torch.Tensor) and latents.ndim == 3):
+            raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
+
+        condition_latents = cls._normalize_ltx2_condition_latents(image_latent)
+        if not condition_latents:
+            return []
+        if len(condition_latents) > 2:
+            raise ValueError(
+                "LTX-2 TI2V currently supports at most two conditioning images."
+            )
+
+        for cond in condition_latents:
+            if not (isinstance(cond, torch.Tensor) and cond.ndim == 3):
+                raise ValueError(
+                    "Expected LTX-2 conditioning latents to be packed tensors [B, S, D]."
+                )
+            if int(cond.shape[1]) < int(num_img_tokens):
+                raise ValueError(
+                    "LTX-2 conditioning latent is shorter than one frame token span."
+                )
+
+        did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
+        if not did_sp_shard:
+            if int(latents.shape[1]) < int(num_img_tokens):
+                raise ValueError(
+                    "LTX-2 latent sequence is shorter than one conditioning frame."
+                )
+            if len(condition_latents) == 1:
+                return [(0, condition_latents[0])]
+            return [
+                (0, condition_latents[0]),
+                (int(latents.shape[1]) - int(num_img_tokens), condition_latents[1]),
+            ]
+
+        tokens_per_frame = int(getattr(batch, "sp_video_tokens_per_frame", 0))
+        if tokens_per_frame <= 0:
+            raise ValueError(
+                "SP-sharded LTX-2 TI2V requires batch.sp_video_tokens_per_frame."
+            )
+        if int(num_img_tokens) != int(tokens_per_frame):
+            raise ValueError(
+                "LTX-2 conditioning token count must match one latent frame when using SP."
+            )
+
+        raw_shape = getattr(batch, "raw_latent_shape", None)
+        if raw_shape is None:
+            raise ValueError("SP-sharded LTX-2 TI2V requires batch.raw_latent_shape.")
+        global_seq_len = int(raw_shape[1])
+        if global_seq_len % tokens_per_frame != 0:
+            raise ValueError(
+                "SP-sharded LTX-2 TI2V expected raw seq_len divisible by tokens_per_frame."
+            )
+
+        global_num_frames = global_seq_len // tokens_per_frame
+        local_start_frame = int(getattr(batch, "sp_video_start_frame", 0))
+        local_num_frames = int(getattr(batch, "sp_video_latent_num_frames", 0))
+        local_end_frame = local_start_frame + local_num_frames
+
+        spans: list[tuple[int, torch.Tensor]] = []
+        if local_start_frame == 0:
+            spans.append((0, condition_latents[0]))
+
+        if len(condition_latents) == 2:
+            last_global_frame = global_num_frames - 1
+            if local_start_frame <= last_global_frame < local_end_frame:
+                local_last_frame = last_global_frame - local_start_frame
+                spans.append(
+                    (local_last_frame * tokens_per_frame, condition_latents[1])
+                )
+
+        return spans
+
+    @classmethod
+    def _prepare_ltx2_ti2v_clean_state(
+        cls,
+        batch: Req,
+        latents: torch.Tensor,
+        image_latent: torch.Tensor | list[torch.Tensor] | None,
         num_img_tokens: int,
         zero_clean_latent: bool,
         clean_latent_background: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         latents = latents.clone()
-        conditioned = image_latent[:, :num_img_tokens, :].to(
-            device=latents.device, dtype=latents.dtype
-        )
-        latents[:, :num_img_tokens, :] = conditioned
         denoise_mask = torch.ones(
             (latents.shape[0], latents.shape[1], 1),
             device=latents.device,
             dtype=torch.float32,
         )
-        denoise_mask[:, :num_img_tokens, :] = 0.0
         if clean_latent_background is not None:
             clean_latent = (
                 clean_latent_background.detach()
@@ -544,7 +631,24 @@ class LTX2DenoisingStage(DenoisingStage):
             clean_latent = torch.zeros_like(latents)
         else:
             clean_latent = latents.detach().clone()
-        clean_latent[:, :num_img_tokens, :] = conditioned
+
+        spans = cls._get_ltx2_condition_spans(
+            batch=batch,
+            latents=latents,
+            image_latent=image_latent,
+            num_img_tokens=num_img_tokens,
+        )
+        for start, cond in spans:
+            stop = int(start) + int(num_img_tokens)
+            conditioned = cls._repeat_batch_dim(
+                cond[:, :num_img_tokens, :].to(
+                    device=latents.device, dtype=latents.dtype
+                ),
+                int(latents.shape[0]),
+            )
+            latents[:, start:stop, :] = conditioned
+            denoise_mask[:, start:stop, :] = 0.0
+            clean_latent[:, start:stop, :] = conditioned
         return latents, denoise_mask, clean_latent
 
     @staticmethod
@@ -915,23 +1019,6 @@ class LTX2DenoisingStage(DenoisingStage):
         super()._preprocess_sp_latents(batch, server_args)
         batch.image_latent = saved
 
-    @staticmethod
-    def _should_apply_ltx2_ti2v(batch: Req) -> bool:
-        """True if we have an image-latent token prefix to condition with.
-
-        SP note: when token latents are time-sharded, only the rank that owns the
-        *global* first latent frame should apply TI2V conditioning (rank with start_frame==0).
-        """
-        if (
-            batch.image_latent is None
-            or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0
-        ):
-            return False
-        did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
-        if not did_sp_shard:
-            return True
-        return int(getattr(batch, "sp_video_start_frame", 0)) == 0
-
     @staticmethod
     def _should_use_native_hq_res2s_sde_noise(server_args: ServerArgs) -> bool:
         return server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
@@ -998,8 +1085,6 @@ class LTX2DenoisingStage(DenoisingStage):
         # Video and audio keep separate scheduler state throughout the denoising loop.
         ctx.audio_scheduler = copy.deepcopy(self.scheduler)
 
-        do_ti2v = self._should_apply_ltx2_ti2v(batch)
-
         if ctx.use_ltx23_legacy_one_stage:
             batch.ltx23_audio_replicated_for_sp = False
             batch.did_sp_shard_audio_latents = False
@@ -1038,6 +1123,13 @@ class LTX2DenoisingStage(DenoisingStage):
             batch.width
             // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
         )
+        ti2v_spans = self._get_ltx2_condition_spans(
+            batch=batch,
+            latents=ctx.latents,
+            image_latent=batch.image_latent,
+            num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
+        )
+        do_ti2v = bool(ti2v_spans)
         if do_ti2v:
             if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3):
                 raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
@@ -1052,6 +1144,7 @@ class LTX2DenoisingStage(DenoisingStage):
             # Keep conditioned tokens clean and reuse the mask during every step update.
             ctx.latents, ctx.denoise_mask, ctx.clean_latent = (
                 self._prepare_ltx2_ti2v_clean_state(
+                    batch=batch,
                     latents=ctx.latents,
                     image_latent=batch.image_latent,
                     num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py

@@ -84,8 +84,8 @@ class LTX2UpsampleStage(PipelineStage):
             device=device, dtype=latents.dtype
         )
         latents = self.spatial_upsampler(latents)
-        if server_args.vae_cpu_offload:
-            self.spatial_upsampler = self.spatial_upsampler.to("cpu")
+        # Keep the small spatial upsampler resident after warmup; moving it
+        # every request dominates the measured two-stage upsample latency.
         latents = (latents - vae_mean) / vae_std
         return latents
 
@@ -118,18 +118,31 @@ class LTX2UpsampleStage(PipelineStage):
         )
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        delay_stage2_prefetch = False
+        if self.pipeline is not None:
+            prepare_upsample = getattr(
+                self.pipeline, "prepare_ltx2_upsample_after_stage1", None
+            )
+            if callable(prepare_upsample):
+                delay_stage2_prefetch = prepare_upsample()
         prefetch_stage2 = (
             getattr(self.pipeline, "prefetch_ltx2_stage2_after_stage1", None)
             if self.pipeline is not None
             else None
         )
-        if callable(prefetch_stage2):
+        if callable(prefetch_stage2) and not delay_stage2_prefetch:
             prefetch_stage2()
 
         device = get_local_torch_device()
         latents = self._upsample_video_latents(batch.latents, server_args, device)
+        if callable(prefetch_stage2) and delay_stage2_prefetch:
+            prefetch_stage2()
         logger.info("Upsampled video latents: %s", list(latents.shape))
         self._restore_full_resolution(batch)
+        batch.image_latent = None
+        batch.ltx2_num_image_tokens = 0
+        batch.did_sp_shard_latents = False
+        batch.did_sp_shard_audio_latents = False
         self._pack_video_latents(batch, latents, server_args)
         logger.info(
             "Packed video latents for Stage 2: %s (resolution %dx%d)",

python/sglang/multimodal_gen/test/server/accuracy_config.py

@@ -142,17 +142,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
             "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
         ),
     },
-    "wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
-        ComponentType.VAE: ComponentSkip(
-            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
-        ),
-        ComponentType.TRANSFORMER: ComponentSkip(
-            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
-        ),
-        ComponentType.TEXT_ENCODER: ComponentSkip(
-            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
-        ),
-    },
     "wan2_1_t2v_1.3b_teacache_enabled": {
         ComponentType.VAE: ComponentSkip(
             "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"

python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py

@@ -31,7 +31,6 @@ ACCURACY_ONE_GPU_CASE_IDS = (
     "flux_2_image_t2i_upscaling_4x",
     "mova_360p_1gpu",
     "wan2_1_t2v_1.3b",
-    "wan2_1_t2v_1.3b_text_encoder_cpu_offload",
     "wan2_1_t2v_1.3b_teacache_enabled",
     "wan2_1_t2v_1.3b_frame_interp_2x",
     "wan2_1_t2v_1.3b_upscaling_4x",

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -175,14 +175,6 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
         ),
         T2V_sampling_params,
     ),
-    DiffusionTestCase(
-        "wan2_1_t2v_1.3b_text_encoder_cpu_offload",
-        DiffusionServerArgs(
-            model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
-            text_encoder_cpu_offload=True,
-        ),
-        T2V_sampling_params,
-    ),
     # TeaCache acceleration test for Wan video model
     DiffusionTestCase(
         "wan2_1_t2v_1.3b_teacache_enabled",

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -951,72 +951,6 @@
             "expected_median_denoise_ms": 145.57,
             "estimated_full_test_time_s": 129.3
         },
-        "wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
-            "stages_ms": {
-                "DecodingStage": 675.91,
-                "TextEncodingStage": 1072.93,
-                "TimestepPreparationStage": 2.43,
-                "LatentPreparationStage": 0.14,
-                "InputValidationStage": 0.07,
-                "DenoisingStage": 7221.86
-            },
-            "denoise_step_ms": {
-                "0": 128.64,
-                "1": 108.46,
-                "2": 140.93,
-                "3": 140.8,
-                "4": 141.21,
-                "5": 141.84,
-                "6": 141.56,
-                "7": 142.35,
-                "8": 142.08,
-                "9": 141.68,
-                "10": 141.69,
-                "11": 141.46,
-                "12": 141.39,
-                "13": 141.36,
-                "14": 141.65,
-                "15": 141.55,
-                "16": 142.02,
-                "17": 141.53,
-                "18": 140.98,
-                "19": 142.4,
-                "20": 141.84,
-                "21": 141.3,
-                "22": 141.41,
-                "23": 141.47,
-                "24": 141.78,
-                "25": 141.6,
-                "26": 142.19,
-                "27": 141.09,
-                "28": 141.2,
-                "29": 141.22,
-                "30": 141.2,
-                "31": 141.23,
-                "32": 141.41,
-                "33": 141.5,
-                "34": 141.56,
-                "35": 141.51,
-                "36": 141.25,
-                "37": 141.49,
-                "38": 141.56,
-                "39": 141.52,
-                "40": 141.17,
-                "41": 141.83,
-                "42": 141.72,
-                "43": 142.31,
-                "44": 141.56,
-                "45": 141.91,
-                "46": 141.93,
-                "47": 141.9,
-                "48": 141.52,
-                "49": 141.34
-            },
-            "expected_e2e_ms": 9296.17,
-            "expected_avg_denoise_ms": 144.33,
-            "expected_median_denoise_ms": 144.84,
-            "estimated_full_test_time_s": 129.3
-        },
         "wan2_1_t2v_1.3b_cfg_parallel": {
             "stages_ms": {
                 "LatentPreparationStage": 0.29,