Update hf_gradio_app.py

hf_gradio_app.py

@@ -65,24 +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
@@ -104,6 +102,7 @@ def process_audio(file_path, temp_dir):
 #@torch.inference_mode()
 @spaces.GPU(duration=200)
 def generate(input_video, input_audio, seed, progress=gr.Progress(track_tqdm=True)):
+    pipeline.to(device=device, dtype=weight_dtype)
 
     is_shared_ui = True if "fffiloni/MEMO" in os.environ['SPACE_ID'] else False
     temp_dir = None