debug

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 
-# from src.audio2vid import audio2video
-# from src.vid2vid import video2video
+from src.audio2vid import audio2video
+from src.vid2vid import video2video
 
-from src.create_modules import Processer
+# from src.create_modules import Processer
 
 title = r"""
 <h1>AniPortrait</h1>
@@ -13,7 +13,7 @@ description = r"""
 <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/Zejun-Yang/AniPortrait' target='_blank'><b>AniPortrait: Audio-Driven Synthesis of Photorealistic Portrait Animations</b></a>.<br>
 """
 
-main_processer = Processer()
+# main_processer = Processer()
 
 with gr.Blocks() as demo:
     
@@ -77,13 +77,13 @@ with gr.Blocks() as demo:
         )
             
     a2v_botton.click(
-        fn=main_processer.audio2video,
+        fn=audio2video,
         inputs=[a2v_input_audio, a2v_ref_img, a2v_headpose_video,
                 a2v_size_slider, a2v_step_slider, a2v_length, a2v_seed], 
         outputs=[a2v_output_video, a2v_ref_img]
     )
     v2v_botton.click(
-        fn=main_processer.video2video,
+        fn=video2video,
         inputs=[v2v_ref_img, v2v_source_video,
                 v2v_size_slider, v2v_step_slider, v2v_length, v2v_seed], 
         outputs=[v2v_output_video, v2v_ref_img]

src/audio2vid.py

@@ -9,27 +9,26 @@ import spaces
 from scipy.spatial.transform import Rotation as R
 from scipy.interpolate import interp1d
 
-# from diffusers import AutoencoderKL, DDIMScheduler
-# from einops import repeat
+from diffusers import AutoencoderKL, DDIMScheduler
+from einops import repeat
 from omegaconf import OmegaConf
 from PIL import Image
 from torchvision import transforms
-# from transformers import CLIPVisionModelWithProjection
+from transformers import CLIPVisionModelWithProjection
 
 
-# from src.models.pose_guider import PoseGuider
-# from src.models.unet_2d_condition import UNet2DConditionModel
-# from src.models.unet_3d import UNet3DConditionModel
-# from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
+from src.models.pose_guider import PoseGuider
+from src.models.unet_2d_condition import UNet2DConditionModel
+from src.models.unet_3d import UNet3DConditionModel
+from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
 from src.utils.util import save_videos_grid
 
-# from src.audio_models.model import Audio2MeshModel
+from src.audio_models.model import Audio2MeshModel
 from src.utils.audio_util import prepare_audio_feature
-# from src.utils.mp_utils  import LMKExtractor
-# from src.utils.draw_util import FaceMeshVisualizer
+from src.utils.mp_utils  import LMKExtractor
+from src.utils.draw_util import FaceMeshVisualizer
 from src.utils.pose_util import project_points
 from src.utils.crop_face_single import crop_face
-from src.create_modules import lmk_extractor, vis, a2m_model, pipe
 
 
 def matrix_to_euler_and_translation(matrix):
@@ -51,7 +50,7 @@ def smooth_pose_seq(pose_seq, window_size=5):
     return smoothed_pose_seq
 
 def get_headpose_temp(input_video):
-    # lmk_extractor = LMKExtractor()
+    lmk_extractor = LMKExtractor()
     cap = cv2.VideoCapture(input_video)
 
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -100,70 +99,70 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
 
     config = OmegaConf.load('./configs/prompts/animation_audio.yaml')
 
-    # if config.weight_dtype == "fp16":
-    #     weight_dtype = torch.float16
-    # else:
-    #     weight_dtype = torch.float32
+    if config.weight_dtype == "fp16":
+        weight_dtype = torch.float16
+    else:
+        weight_dtype = torch.float32
         
     audio_infer_config = OmegaConf.load(config.audio_inference_config)
-    # # prepare model
-    # a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
-    # a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt']), strict=False)
-    # a2m_model.cuda().eval()
+    # prepare model
+    a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
+    a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt'], map_location="cpu"), strict=False)
+    a2m_model.cuda().eval()
 
-    # vae = AutoencoderKL.from_pretrained(
-    #     config.pretrained_vae_path,
-    # ).to("cuda", dtype=weight_dtype)
+    vae = AutoencoderKL.from_pretrained(
+        config.pretrained_vae_path,
+    ).to("cuda", dtype=weight_dtype)
 
-    # reference_unet = UNet2DConditionModel.from_pretrained(
-    #     config.pretrained_base_model_path,
-    #     subfolder="unet",
-    # ).to(dtype=weight_dtype, device="cuda")
+    reference_unet = UNet2DConditionModel.from_pretrained(
+        config.pretrained_base_model_path,
+        subfolder="unet",
+    ).to(dtype=weight_dtype, device="cuda")
 
-    # inference_config_path = config.inference_config
-    # infer_config = OmegaConf.load(inference_config_path)
-    # denoising_unet = UNet3DConditionModel.from_pretrained_2d(
-    #     config.pretrained_base_model_path,
-    #     config.motion_module_path,
-    #     subfolder="unet",
-    #     unet_additional_kwargs=infer_config.unet_additional_kwargs,
-    # ).to(dtype=weight_dtype, device="cuda")
+    inference_config_path = config.inference_config
+    infer_config = OmegaConf.load(inference_config_path)
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        config.motion_module_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(dtype=weight_dtype, device="cuda")
 
 
-    # pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
+    pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
 
-    # image_enc = CLIPVisionModelWithProjection.from_pretrained(
-    #     config.image_encoder_path
-    # ).to(dtype=weight_dtype, device="cuda")
+    image_enc = CLIPVisionModelWithProjection.from_pretrained(
+        config.image_encoder_path
+    ).to(dtype=weight_dtype, device="cuda")
 
-    # sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
-    # scheduler = DDIMScheduler(**sched_kwargs)
+    sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
+    scheduler = DDIMScheduler(**sched_kwargs)
 
     generator = torch.manual_seed(seed)
 
     width, height = size, size
 
-    # # load pretrained weights
-    # denoising_unet.load_state_dict(
-    #     torch.load(config.denoising_unet_path, map_location="cpu"),
-    #     strict=False,
-    # )
-    # reference_unet.load_state_dict(
-    #     torch.load(config.reference_unet_path, map_location="cpu"),
-    # )
-    # pose_guider.load_state_dict(
-    #     torch.load(config.pose_guider_path, map_location="cpu"),
-    # )
-
-    # pipe = Pose2VideoPipeline(
-    #     vae=vae,
-    #     image_encoder=image_enc,
-    #     reference_unet=reference_unet,
-    #     denoising_unet=denoising_unet,
-    #     pose_guider=pose_guider,
-    #     scheduler=scheduler,
-    # )
-    # pipe = pipe.to("cuda", dtype=weight_dtype)
+    # load pretrained weights
+    denoising_unet.load_state_dict(
+        torch.load(config.denoising_unet_path, map_location="cpu"),
+        strict=False,
+    )
+    reference_unet.load_state_dict(
+        torch.load(config.reference_unet_path, map_location="cpu"),
+    )
+    pose_guider.load_state_dict(
+        torch.load(config.pose_guider_path, map_location="cpu"),
+    )
+
+    pipe = Pose2VideoPipeline(
+        vae=vae,
+        image_encoder=image_enc,
+        reference_unet=reference_unet,
+        denoising_unet=denoising_unet,
+        pose_guider=pose_guider,
+        scheduler=scheduler,
+    )
+    pipe = pipe.to("cuda", dtype=weight_dtype)
 
     date_str = datetime.now().strftime("%Y%m%d")
     time_str = datetime.now().strftime("%H%M")
@@ -172,8 +171,8 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     save_dir = Path(f"output/{date_str}/{save_dir_name}")
     save_dir.mkdir(exist_ok=True, parents=True)
 
-    # lmk_extractor = LMKExtractor()
-    # vis = FaceMeshVisualizer(forehead_edge=False)
+    lmk_extractor = LMKExtractor()
+    vis = FaceMeshVisualizer(forehead_edge=False)
 
     ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
     ref_image_np = crop_face(ref_image_np, lmk_extractor)

src/vid2vid.py

@@ -7,89 +7,88 @@ import numpy as np
 import cv2
 import torch
 import spaces
-# from diffusers import AutoencoderKL, DDIMScheduler
-# from einops import repeat
-# from omegaconf import OmegaConf
+from diffusers import AutoencoderKL, DDIMScheduler
+from einops import repeat
+from omegaconf import OmegaConf
 from PIL import Image
 from torchvision import transforms
-# from transformers import CLIPVisionModelWithProjection
+from transformers import CLIPVisionModelWithProjection
 
-# from src.models.pose_guider import PoseGuider
-# from src.models.unet_2d_condition import UNet2DConditionModel
-# from src.models.unet_3d import UNet3DConditionModel
-# from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
+from src.models.pose_guider import PoseGuider
+from src.models.unet_2d_condition import UNet2DConditionModel
+from src.models.unet_3d import UNet3DConditionModel
+from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
 from src.utils.util import get_fps, read_frames, save_videos_grid
 
-# from src.utils.mp_utils  import LMKExtractor
-# from src.utils.draw_util import FaceMeshVisualizer
+from src.utils.mp_utils  import LMKExtractor
+from src.utils.draw_util import FaceMeshVisualizer
 from src.utils.pose_util import project_points_with_trans, matrix_to_euler_and_translation, euler_and_translation_to_matrix
 from src.audio2vid import smooth_pose_seq
 from src.utils.crop_face_single import crop_face
-from src.create_modules import lmk_extractor, vis, pipe
 
 @spaces.GPU
 def video2video(ref_img, source_video, size=512, steps=25, length=150, seed=42):
     cfg = 3.5
     
-    # config = OmegaConf.load('./configs/prompts/animation_facereenac.yaml')
+    config = OmegaConf.load('./configs/prompts/animation_facereenac.yaml')
 
-    # if config.weight_dtype == "fp16":
-    #     weight_dtype = torch.float16
-    # else:
-    #     weight_dtype = torch.float32
+    if config.weight_dtype == "fp16":
+        weight_dtype = torch.float16
+    else:
+        weight_dtype = torch.float32
 
-    # vae = AutoencoderKL.from_pretrained(
-    #     config.pretrained_vae_path,
-    # ).to("cuda", dtype=weight_dtype)
+    vae = AutoencoderKL.from_pretrained(
+        config.pretrained_vae_path,
+    ).to("cuda", dtype=weight_dtype)
 
-    # reference_unet = UNet2DConditionModel.from_pretrained(
-    #     config.pretrained_base_model_path,
-    #     subfolder="unet",
-    # ).to(dtype=weight_dtype, device="cuda")
+    reference_unet = UNet2DConditionModel.from_pretrained(
+        config.pretrained_base_model_path,
+        subfolder="unet",
+    ).to(dtype=weight_dtype, device="cuda")
 
-    # inference_config_path = config.inference_config
-    # infer_config = OmegaConf.load(inference_config_path)
-    # denoising_unet = UNet3DConditionModel.from_pretrained_2d(
-    #     config.pretrained_base_model_path,
-    #     config.motion_module_path,
-    #     subfolder="unet",
-    #     unet_additional_kwargs=infer_config.unet_additional_kwargs,
-    # ).to(dtype=weight_dtype, device="cuda")
+    inference_config_path = config.inference_config
+    infer_config = OmegaConf.load(inference_config_path)
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        config.motion_module_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(dtype=weight_dtype, device="cuda")
 
-    # pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
+    pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
 
-    # image_enc = CLIPVisionModelWithProjection.from_pretrained(
-    #     config.image_encoder_path
-    # ).to(dtype=weight_dtype, device="cuda")
+    image_enc = CLIPVisionModelWithProjection.from_pretrained(
+        config.image_encoder_path
+    ).to(dtype=weight_dtype, device="cuda")
 
-    # sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
-    # scheduler = DDIMScheduler(**sched_kwargs)
+    sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
+    scheduler = DDIMScheduler(**sched_kwargs)
 
     generator = torch.manual_seed(seed)
 
     width, height = size, size
 
-    # # load pretrained weights
-    # denoising_unet.load_state_dict(
-    #     torch.load(config.denoising_unet_path, map_location="cpu"),
-    #     strict=False,
-    # )
-    # reference_unet.load_state_dict(
-    #     torch.load(config.reference_unet_path, map_location="cpu"),
-    # )
-    # pose_guider.load_state_dict(
-    #     torch.load(config.pose_guider_path, map_location="cpu"),
-    # )
-
-    # pipe = Pose2VideoPipeline(
-    #     vae=vae,
-    #     image_encoder=image_enc,
-    #     reference_unet=reference_unet,
-    #     denoising_unet=denoising_unet,
-    #     pose_guider=pose_guider,
-    #     scheduler=scheduler,
-    # )
-    # pipe = pipe.to("cuda", dtype=weight_dtype)
+    # load pretrained weights
+    denoising_unet.load_state_dict(
+        torch.load(config.denoising_unet_path, map_location="cpu"),
+        strict=False,
+    )
+    reference_unet.load_state_dict(
+        torch.load(config.reference_unet_path, map_location="cpu"),
+    )
+    pose_guider.load_state_dict(
+        torch.load(config.pose_guider_path, map_location="cpu"),
+    )
+
+    pipe = Pose2VideoPipeline(
+        vae=vae,
+        image_encoder=image_enc,
+        reference_unet=reference_unet,
+        denoising_unet=denoising_unet,
+        pose_guider=pose_guider,
+        scheduler=scheduler,
+    )
+    pipe = pipe.to("cuda", dtype=weight_dtype)
 
     date_str = datetime.now().strftime("%Y%m%d")
     time_str = datetime.now().strftime("%H%M")
@@ -99,11 +98,9 @@ def video2video(ref_img, source_video, size=512, steps=25, length=150, seed=42):
     save_dir.mkdir(exist_ok=True, parents=True)
 
 
-    # lmk_extractor = LMKExtractor()
-    # vis = FaceMeshVisualizer(forehead_edge=False)
+    lmk_extractor = LMKExtractor()
+    vis = FaceMeshVisualizer(forehead_edge=False)
     
-
-
     ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
     ref_image_np = crop_face(ref_image_np, lmk_extractor)
     if ref_image_np is None: