Update hf_gradio_app.py

hf_gradio_app.py

@@ -65,22 +65,22 @@ from memo.utils.vision_utils import preprocess_image, tensor_to_video
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 weight_dtype = torch.bfloat16
 
-#with torch.inference_mode():
-vae = AutoencoderKL.from_pretrained("./checkpoints/vae").to(device=device, dtype=weight_dtype)
-reference_net = UNet2DConditionModel.from_pretrained("./checkpoints", subfolder="reference_net", use_safetensors=True)
-diffusion_net = UNet3DConditionModel.from_pretrained("./checkpoints", subfolder="diffusion_net", use_safetensors=True)
-image_proj = ImageProjModel.from_pretrained("./checkpoints", subfolder="image_proj", use_safetensors=True)
-audio_proj = AudioProjModel.from_pretrained("./checkpoints", subfolder="audio_proj", use_safetensors=True)
-vae.requires_grad_(False).eval()
-reference_net.requires_grad_(False).eval()
-diffusion_net.requires_grad_(False).eval()
-image_proj.requires_grad_(False).eval()
-audio_proj.requires_grad_(False).eval()
-#reference_net.enable_xformers_memory_efficient_attention()
-#diffusion_net.enable_xformers_memory_efficient_attention()
-noise_scheduler = FlowMatchEulerDiscreteScheduler()
-pipeline = VideoPipeline(vae=vae, reference_net=reference_net, diffusion_net=diffusion_net, scheduler=noise_scheduler, image_proj=image_proj)
-#pipeline.to(device=device, dtype=weight_dtype)
+with torch.inference_mode():
+    vae = AutoencoderKL.from_pretrained("./checkpoints/vae").to(device=device, dtype=weight_dtype)
+    reference_net = UNet2DConditionModel.from_pretrained("./checkpoints", subfolder="reference_net", use_safetensors=True)
+    diffusion_net = UNet3DConditionModel.from_pretrained("./checkpoints", subfolder="diffusion_net", use_safetensors=True)
+    image_proj = ImageProjModel.from_pretrained("./checkpoints", subfolder="image_proj", use_safetensors=True)
+    audio_proj = AudioProjModel.from_pretrained("./checkpoints", subfolder="audio_proj", use_safetensors=True)
+    vae.requires_grad_(False).eval()
+    reference_net.requires_grad_(False).eval()
+    diffusion_net.requires_grad_(False).eval()
+    image_proj.requires_grad_(False).eval()
+    audio_proj.requires_grad_(False).eval()
+    #reference_net.enable_xformers_memory_efficient_attention()
+    #diffusion_net.enable_xformers_memory_efficient_attention()
+    noise_scheduler = FlowMatchEulerDiscreteScheduler()
+    pipeline = VideoPipeline(vae=vae, reference_net=reference_net, diffusion_net=diffusion_net, scheduler=noise_scheduler, image_proj=image_proj)
+    #pipeline.to(device=device, dtype=weight_dtype)
 
 def process_audio(file_path, temp_dir):
     # Load the audio file