Update app.py

app.py

@@ -39,7 +39,7 @@ from pipeline import CausalInferencePipeline
 from demo_utils.constant import ZERO_VAE_CACHE
 from demo_utils.vae_block3 import VAEDecoderWrapper
 from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
-from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
+# from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
 
 # --- Argument Parsing ---
 parser = argparse.ArgumentParser(description="Gradio Demo for Self-Forcing with FastRTC")
@@ -51,12 +51,8 @@ parser.add_argument('--share', action='store_true', help="Create a public Gradio
 parser.add_argument('--trt', action='store_true', help="Use TensorRT optimized VAE decoder.")
 args = parser.parse_args()
 
-# --- Global Setup & Model Loading ---
-print(f"CUDA device: {gpu}")
-print(f'Initial Free VRAM: {get_cuda_free_memory_gb(gpu):.2f} GB')
-LOW_MEMORY = get_cuda_free_memory_gb(gpu) < 40
+gpu = "cuda"
 
-# Load configs
 try:
     config = OmegaConf.load(args.config_path)
     default_config = OmegaConf.load("configs/default_config.yaml")
@@ -77,18 +73,12 @@ except FileNotFoundError as e:
     print(f"Error loading checkpoint: {e}\nPlease ensure the checkpoint '{args.checkpoint_path}' exists.")
     exit(1)
 
-# Prepare models for inference
 text_encoder.eval().to(dtype=torch.bfloat16).requires_grad_(False)
 transformer.eval().to(dtype=torch.float16).requires_grad_(False)
 
-if LOW_MEMORY:
-    print("Low memory mode enabled. Using dynamic model swapping.")
-    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
-else:
-    text_encoder.to(gpu)
+text_encoder.to(gpu)
 transformer.to(gpu)
 
-# --- VAE Decoder Management ---
 APP_STATE = {
     "torch_compile_applied": False,
     "fp8_applied": False,