vae update

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")

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 = {

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)

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
 

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.")