Update app.py

app.py

@@ -45,7 +45,7 @@ hf_hub_download(repo_id="wileewang/TransPixar", filename="cogvideox_rgba_lora.sa
 # snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife")
 
 pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5B", torch_dtype=torch.bfloat16)
-# pipe.enable_sequential_cpu_offload()
+pipe.enable_sequential_cpu_offload()
 pipe.vae.enable_slicing()
 pipe.vae.enable_tiling()
 pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
@@ -57,7 +57,7 @@ seq_length = 2 * (
 prepare_for_rgba_inference(
     pipe.transformer,
     rgba_weights_path="model_cogvideox_rgba_lora/cogvideox_rgba_lora.safetensors",
-    device="cuda",
+    device=device,
     dtype=torch.bfloat16,
     text_length=226,
     seq_length=seq_length, # this is for the creation of attention mask.
@@ -308,6 +308,8 @@ with gr.Blocks() as demo:
         alpha_video_update = gr.update(visible=True, value=alpha_video_path)
         seed_update = gr.update(visible=True, value=seed)
 
+        pipe.to("cpu")
+
         return rgb_video_path, alpha_video_path, rgb_video_update, alpha_video_update, seed_update