From 3bd2a885e46ec1e2b099a83edbb599607a4f0c96 Mon Sep 17 00:00:00 2001
From: Will Lin <wlsaidhi@gmail.com>
Date: Fri, 6 Feb 2026 15:41:09 -0800
Subject: [PATCH] update

---
 .../inference/basic/basic_ltx2_upscale.py     |  6 ++--
 fastvideo/pipelines/stages/ltx2_refine.py     | 30 +++++++++++--------
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/examples/inference/basic/basic_ltx2_upscale.py b/examples/inference/basic/basic_ltx2_upscale.py
index b4b06ca54..70acbe572 100644
--- a/examples/inference/basic/basic_ltx2_upscale.py
+++ b/examples/inference/basic/basic_ltx2_upscale.py
@@ -34,12 +34,12 @@ def main() -> None:
     generator.generate_video(
         prompt=PROMPT,
         output_path=OUTPUT_PATH,
-        height=1024,
-        width=1792,
+        height=960,
+        width=1664,
         num_frames=81,
         fps=24,
         seed=10,
-        # Distilled schedule in LTX-2 uses 8 steps; setting this keeps parity.
+        # DistilledPipeline uses the 8-step distilled schedule without CFG.
         num_inference_steps=8,
         guidance_scale=1.0,
         save_video=True,
diff --git a/fastvideo/pipelines/stages/ltx2_refine.py b/fastvideo/pipelines/stages/ltx2_refine.py
index 6fa5d738c..71f3f31e2 100644
--- a/fastvideo/pipelines/stages/ltx2_refine.py
+++ b/fastvideo/pipelines/stages/ltx2_refine.py
@@ -115,20 +115,22 @@ class LTX2UpsampleStage(PipelineStage):
             raise ValueError(
                 "Latents must be available before LTX-2 upsample stage.")
 
-        # Ensure upsampler matches latent dtype/device to avoid conv mismatch.
+        latents = batch.latents
+        orig_dtype = latents.dtype
+        orig_device = latents.device
         if isinstance(self.upsampler, torch.nn.Module):
-            target_dtype = batch.latents.dtype
-            target_device = batch.latents.device
             first_param = next(self.upsampler.parameters(), None)
-            if first_param is not None and (first_param.dtype != target_dtype
-                                            or first_param.device
-                                            != target_device):
-                self.upsampler.to(device=target_device, dtype=target_dtype)
-                logger.info(
-                    "[LTX2] Cast upsampler to %s on %s to match latents.",
-                    target_dtype,
-                    target_device,
-                )
+            if first_param is not None:
+                if first_param.device != orig_device:
+                    latents = latents.to(device=first_param.device)
+                if first_param.dtype != latents.dtype:
+                    latents = latents.to(dtype=first_param.dtype)
+                if latents.dtype != orig_dtype or latents.device != orig_device:
+                    logger.info(
+                        "[LTX2] Cast latents to %s on %s for upsampler.",
+                        latents.dtype,
+                        latents.device,
+                    )
 
         target_height = batch.extra.get("ltx2_refine_target_height")
         target_width = batch.extra.get("ltx2_refine_target_width")
@@ -141,7 +143,9 @@ class LTX2UpsampleStage(PipelineStage):
                 "LTX-2 VAE encoder is required for latent upsampling.")
 
         upsampler_module = getattr(self.upsampler, "model", self.upsampler)
-        latents = upsample_video(batch.latents, video_encoder, upsampler_module)
+        latents = upsample_video(latents, video_encoder, upsampler_module)
+        if latents.dtype != orig_dtype or latents.device != orig_device:
+            latents = latents.to(device=orig_device, dtype=orig_dtype)
 
         sigma0 = float(self.sigmas[0]) if self.sigmas else 1.0
         if self.add_noise: