From da901e7368fd6daa616575ea131ef732d1b7c851 Mon Sep 17 00:00:00 2001
From: Shenghai Yuan <140951558+SHYuanBest@users.noreply.github.com>
Date: Thu, 5 Mar 2026 20:16:19 +0800
Subject: [PATCH] Expand _skip_layerwise_casting_patterns list

---
 .../diffusers_version/transformer_helios_diffusers.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/helios/diffusers_version/transformer_helios_diffusers.py b/helios/diffusers_version/transformer_helios_diffusers.py
index d0761dc..bd3d5c2 100644
--- a/helios/diffusers_version/transformer_helios_diffusers.py
+++ b/helios/diffusers_version/transformer_helios_diffusers.py
@@ -536,7 +536,14 @@ class HeliosTransformer3DModel(
     """
 
     _supports_gradient_checkpointing = True
-    _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"]
+    _skip_layerwise_casting_patterns = [
+        "patch_embedding",
+        "patch_short",
+        "patch_mid",
+        "patch_long",
+        "condition_embedder",
+        "norm",
+    ]
     _no_split_modules = ["HeliosTransformerBlock", "HeliosOutputNorm"]
     _keep_in_fp32_modules = [
         "time_embedder",
@@ -808,4 +815,4 @@ class HeliosTransformer3DModel(
         if not return_dict:
             return (output,)
 
-        return Transformer2DModelOutput(sample=output)
\ No newline at end of file
+        return Transformer2DModelOutput(sample=output)