Set device for encode (#420)

Author: kevin314

fastvideo/v1/models/vaes/common.py

@@ -39,6 +39,9 @@ def __init__(self, config: VAEConfig, **kwargs) -> None:
         self.use_temporal_tiling = config.use_temporal_tiling
         self.use_parallel_tiling = config.use_parallel_tiling
 
+    def to(self, device) -> 'ParallelTiledVAE':
+        return self
+
     @property
     def temporal_compression_ratio(self) -> int:
         return cast(int, self.config.temporal_compression_ratio)

fastvideo/v1/pipelines/stages/decoding.py

@@ -7,6 +7,7 @@
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
@@ -23,7 +24,7 @@ class DecodingStage(PipelineStage):
     """
 
     def __init__(self, vae) -> None:
-        self.vae = vae
+        self.vae: ParallelTiledVAE = vae
 
     def forward(
         self,

fastvideo/v1/pipelines/stages/encoding.py

@@ -46,6 +46,8 @@ def forward(
         Returns:
             The batch with encoded outputs.
         """
+        self.vae = self.vae.to(fastvideo_args.device)
+
         image_path = batch.image_path
         # TODO(will): remove this once we add input/output validation for stages
         if image_path is None: