diff --git a/fastvideo/v1/layers/activation.py b/fastvideo/v1/layers/activation.py
index dc7db7d22..6217ff1a5 100644
--- a/fastvideo/v1/layers/activation.py
+++ b/fastvideo/v1/layers/activation.py
@@ -11,8 +11,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size)
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.platforms import current_platform
-from vllm.utils import LazyDict
+from fastvideo.v1.platforms import current_platform
 
 
 @CustomOp.register("fatrelu_and_mul")
diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py
index 8baee0883..53c9b0628 100644
--- a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py
+++ b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py
@@ -69,7 +69,8 @@ MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan")
 DATA_TYPE = {"image", "video", "image_video"}
 
 # 3D VAE
-VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"}
+# VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"}
+VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/vae"}
 
 # Text Encoder
 TEXT_ENCODER_PATH = {
diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py
index 761c7cc1f..aa3de3c84 100644
--- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py
+++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py
@@ -35,7 +35,7 @@ def load_vae(
     else:
         vae = AutoencoderKLCausal3D.from_config(config)
 
-    vae_ckpt = Path(vae_path) / "pytorch_model.pt"
+    vae_ckpt = Path("/home/ubuntu/src/FastVideo/data/hunyuan/hunyuan-video-t2v-720p/vae/pytorch_model.pt")
     assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}"
 
     ckpt = torch.load(vae_ckpt, map_location=vae.device)
diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py
index 5a63783cd..b0973853b 100644
--- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py
+++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py
@@ -25,7 +25,7 @@ import torch.nn as nn
 import torch.distributed as dist
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 
-from fastvideo.v1.utils.parallel_states import nccl_info
+from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info
 from fastvideo.v1.distributed.parallel_state import get_sp_group
 from fastvideo.v1.logger import init_logger
 
diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py
index 3c466f55f..6b16b26bd 100644
--- a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py
+++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py
@@ -448,20 +448,20 @@ def get_down_block3d(
         attention_head_dim = num_attention_heads
 
     down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type)
-    if down_block_type == "DownEncoderBlockCausal3D":
-        return DownEncoderBlockCausal3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            dropout=dropout,
-            add_downsample=add_downsample,
-            downsample_stride=downsample_stride,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            downsample_padding=downsample_padding,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-        )
+    # if down_block_type == "DownEncoderBlockCausal3D":
+    return DownEncoderBlockCausal3D(
+        num_layers=num_layers,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        dropout=dropout,
+        add_downsample=add_downsample,
+        downsample_stride=downsample_stride,
+        resnet_eps=resnet_eps,
+        resnet_act_fn=resnet_act_fn,
+        resnet_groups=resnet_groups,
+        downsample_padding=downsample_padding,
+        resnet_time_scale_shift=resnet_time_scale_shift,
+    )
     raise ValueError(f"{down_block_type} does not exist.")
 
 
@@ -502,21 +502,21 @@ def get_up_block3d(
         attention_head_dim = num_attention_heads
 
     up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type)
-    if up_block_type == "UpDecoderBlockCausal3D":
-        return UpDecoderBlockCausal3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            resolution_idx=resolution_idx,
-            dropout=dropout,
-            add_upsample=add_upsample,
-            upsample_scale_factor=upsample_scale_factor,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            temb_channels=temb_channels,
-        )
+    # if up_block_type == "UpDecoderBlockCausal3D":
+    return UpDecoderBlockCausal3D(
+        num_layers=num_layers,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        resolution_idx=resolution_idx,
+        dropout=dropout,
+        add_upsample=add_upsample,
+        upsample_scale_factor=upsample_scale_factor,
+        resnet_eps=resnet_eps,
+        resnet_act_fn=resnet_act_fn,
+        resnet_groups=resnet_groups,
+        resnet_time_scale_shift=resnet_time_scale_shift,
+        temb_channels=temb_channels,
+    )
     raise ValueError(f"{up_block_type} does not exist.")