diff --git a/fastvideo/__init__.py b/fastvideo/__init__.py
index 44bbc0851..e69de29bb 100644
--- a/fastvideo/__init__.py
+++ b/fastvideo/__init__.py
@@ -1,3 +0,0 @@
-from fastvideo.v1.entrypoints.video_generator import VideoGenerator
-
-__all__ = ["VideoGenerator"]
diff --git a/fastvideo/v1/forward_context.py b/fastvideo/v1/forward_context.py
index 6455271fc..e74dd822e 100644
--- a/fastvideo/v1/forward_context.py
+++ b/fastvideo/v1/forward_context.py
@@ -52,7 +52,7 @@ def get_forward_context() -> ForwardContext:
 @contextmanager
 def set_forward_context(current_timestep,
                         attn_metadata,
-                        fastvideo_args: FastVideoArgs = None):
+                        fastvideo_args: Optional[FastVideoArgs] = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     Here we can inject common logic for every model forward pass.
diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py
index 18b66d373..9af18431e 100644
--- a/fastvideo/v1/pipelines/stages/denoising.py
+++ b/fastvideo/v1/pipelines/stages/denoising.py
@@ -193,7 +193,7 @@ class DenoisingStage(PipelineStage):
                             attn_metadata = self.attn_metadata_builder.build(
                                 current_timestep=i,
                                 forward_batch=batch,
-                                inference_args=inference_args,
+                                inference_args=fastvideo_args,
                             )
                             assert attn_metadata is not None, "attn_metadata cannot be None"
                         else:
diff --git a/fastvideo/v1/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py
index c3424ea98..5e1760cfa 100644
--- a/fastvideo/v1/pipelines/stages/latent_preparation.py
+++ b/fastvideo/v1/pipelines/stages/latent_preparation.py
@@ -69,8 +69,8 @@ class LatentPreparationStage(PipelineStage):
         if height is None or width is None:
             raise ValueError("Height and width must be provided")
 
-        assert inference_args.num_channels_latents is not None
-        assert inference_args.vae_scale_factor is not None
+        assert fastvideo_args.num_channels_latents is not None
+        assert fastvideo_args.vae_scale_factor is not None
 
         # Calculate latent shape
         shape = (
diff --git a/fastvideo/v1/sample/v1_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py
index 3f301508a..508a0cc5c 100644
--- a/fastvideo/v1/sample/v1_fastvideo_inference.py
+++ b/fastvideo/v1/sample/v1_fastvideo_inference.py
@@ -27,6 +27,8 @@ def initialize_distributed_and_parallelism(fastvideo_args: FastVideoArgs):
     device_str = f"cuda:{local_rank}"
     fastvideo_args.device_str = device_str
     fastvideo_args.device = torch.device(device_str)
+    assert fastvideo_args.sp_size is not None
+    assert fastvideo_args.tp_size is not None
     initialize_model_parallel(
         sequence_model_parallel_size=fastvideo_args.sp_size,
         tensor_model_parallel_size=fastvideo_args.tp_size,
@@ -41,6 +43,8 @@ def main(fastvideo_args: FastVideoArgs):
         with open(fastvideo_args.prompt_path) as f:
             prompts = [line.strip() for line in f.readlines()]
     else:
+        if fastvideo_args.prompt is None:
+            raise ValueError("prompt or prompt_path is required")
         prompts = [fastvideo_args.prompt]
 
     # Process each prompt
diff --git a/pyproject.toml b/pyproject.toml
index ce8ff8442..95820cd55 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,7 +72,6 @@ exclude = ["assets*", "docker*", "docs", "scripts*"]
 exclude = ["assets*", "docker*", "docs", "scripts*"]
 
 [tool.mypy]
-warn_return_any = true
 warn_unused_configs = true
 ignore_missing_imports = true
 check_untyped_defs = true