Upload 99 files

.gitattributes

@@ -33,3 +33,18 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/audio/2.WAV filter=lfs diff=lfs merge=lfs -text
+assets/audio/3.WAV filter=lfs diff=lfs merge=lfs -text
+assets/audio/4.WAV filter=lfs diff=lfs merge=lfs -text
+assets/image/1.png filter=lfs diff=lfs merge=lfs -text
+assets/image/2.png filter=lfs diff=lfs merge=lfs -text
+assets/image/3.png filter=lfs diff=lfs merge=lfs -text
+assets/image/4.png filter=lfs diff=lfs merge=lfs -text
+assets/image/src1.png filter=lfs diff=lfs merge=lfs -text
+assets/image/src2.png filter=lfs diff=lfs merge=lfs -text
+assets/image/src3.png filter=lfs diff=lfs merge=lfs -text
+assets/image/src4.png filter=lfs diff=lfs merge=lfs -text
+assets/material/demo.png filter=lfs diff=lfs merge=lfs -text
+assets/material/logo.png filter=lfs diff=lfs merge=lfs -text
+assets/material/method.png filter=lfs diff=lfs merge=lfs -text
+assets/material/teaser.png filter=lfs diff=lfs merge=lfs -text

assets/audio/2.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:388d768354bb20f3aaa8327ef3391737c8150d3351fdaa04aa98b57caddc5dfb
+size 3862572

assets/audio/3.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:854f2df163f4fc7bd69f09e8ca31758dedf8a56b191083d5fd4a5b25259b5fe2
+size 1921068

assets/audio/4.WAV

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81d992d96a96829f27aef2d99fa63614b692cc9ed2cbac61e23e9a1dcc402b3a
+size 3809324

assets/test.csv

@@ -0,0 +1,25 @@
+videoid,image,audio,prompt,fps
+8,assets/image/1.png,assets/audio/2.WAV,A person sits cross-legged by a campfire in a forested area.,25
+9,assets/image/2.png,assets/audio/2.WAV,"A person with long blonde hair wearing a green jacket, standing in a forested area during twilight.",25
+10,assets/image/3.png,assets/audio/2.WAV,A person playing guitar by a campfire in a forest.,25
+11,assets/image/4.png,assets/audio/2.WAV,"A person wearing a green jacket stands in a forested area, with sunlight filtering through the trees.",25
+12,assets/image/src1.png,assets/audio/2.WAV,A person sits cross-legged by a campfire in a forest at dusk.,25
+13,assets/image/src2.png,assets/audio/2.WAV,A person in a green jacket stands in a forest at dusk.,25
+14,assets/image/src3.png,assets/audio/2.WAV,A person playing guitar by a campfire in a forest.,25
+15,assets/image/src4.png,assets/audio/2.WAV,"A person in a green jacket stands in a forest, backlit by sunlight.",25
+16,assets/image/1.png,assets/audio/3.WAV,A person sits cross-legged by a campfire in a forested area.,25
+17,assets/image/2.png,assets/audio/3.WAV,"A person with long blonde hair wearing a green jacket, standing in a forested area during twilight.",25
+18,assets/image/3.png,assets/audio/3.WAV,A person playing guitar by a campfire in a forest.,25
+19,assets/image/4.png,assets/audio/3.WAV,"A person wearing a green jacket stands in a forested area, with sunlight filtering through the trees.",25
+20,assets/image/src1.png,assets/audio/3.WAV,A person sits cross-legged by a campfire in a forest at dusk.,25
+21,assets/image/src2.png,assets/audio/3.WAV,A person in a green jacket stands in a forest at dusk.,25
+22,assets/image/src3.png,assets/audio/3.WAV,A person playing guitar by a campfire in a forest.,25
+23,assets/image/src4.png,assets/audio/3.WAV,"A person in a green jacket stands in a forest, backlit by sunlight.",25
+24,assets/image/1.png,assets/audio/4.WAV,A person sits cross-legged by a campfire in a forested area.,25
+25,assets/image/2.png,assets/audio/4.WAV,"A person with long blonde hair wearing a green jacket, standing in a forested area during twilight.",25
+26,assets/image/3.png,assets/audio/4.WAV,A person playing guitar by a campfire in a forest.,25
+27,assets/image/4.png,assets/audio/4.WAV,"A person wearing a green jacket stands in a forested area, with sunlight filtering through the trees.",25
+28,assets/image/src1.png,assets/audio/4.WAV,A person sits cross-legged by a campfire in a forest at dusk.,25
+29,assets/image/src2.png,assets/audio/4.WAV,A person in a green jacket stands in a forest at dusk.,25
+30,assets/image/src3.png,assets/audio/4.WAV,A person playing guitar by a campfire in a forest.,25
+31,assets/image/src4.png,assets/audio/4.WAV,"A person in a green jacket stands in a forest, backlit by sunlight.",25

hymm_gradio/flask_audio.py

@@ -0,0 +1,268 @@
+import os
+import numpy as np
+import torch
+import warnings
+import threading
+import traceback
+import uvicorn
+from fastapi import FastAPI, Body
+from pathlib import Path
+from datetime import datetime
+import torch.distributed as dist
+from hymm_gradio.tool_for_end2end import *
+from hymm_sp.config import parse_args
+from hymm_sp.sample_inference_audio import HunyuanVideoSampler
+
+from hymm_sp.modules.parallel_states import (
+    initialize_distributed,
+    nccl_info,
+)
+
+from transformers import WhisperModel
+from transformers import AutoFeatureExtractor
+from hymm_sp.data_kits.face_align import AlignImage
+
+
+warnings.filterwarnings("ignore")
+MODEL_OUTPUT_PATH = os.environ.get('MODEL_BASE')
+app = FastAPI()
+rlock = threading.RLock()
+
+
+
+@app.api_route('/predict2', methods=['GET', 'POST'])
+def predict(data=Body(...)):
+    is_acquire = False
+    error_info = ""
+    try:
+        is_acquire = rlock.acquire(blocking=False)
+        if is_acquire:
+            res = predict_wrap(data)
+            return res
+    except Exception as e:
+        error_info = traceback.format_exc()
+        print(error_info)
+    finally:
+        if is_acquire:
+            rlock.release()
+    return {"errCode": -1, "info": "broken"}
+
+def predict_wrap(input_dict={}):
+    if nccl_info.sp_size > 1:
+        device = torch.device(f"cuda:{torch.distributed.get_rank()}")
+        rank = local_rank = torch.distributed.get_rank()
+        print(f"sp_size={nccl_info.sp_size}, rank {rank} local_rank {local_rank}")
+    try:
+        print(f"----- rank = {rank}")
+        if rank == 0:
+            input_dict = process_input_dict(input_dict)
+
+            print('------- start to predict -------')
+            # Parse input arguments
+            image_path = input_dict["image_path"]
+            driving_audio_path = input_dict["audio_path"]
+
+            prompt = input_dict["prompt"]
+
+            save_fps = input_dict.get("save_fps", 25)
+
+
+            ret_dict = None
+            if image_path is None or driving_audio_path is None:
+                ret_dict = {
+                    "errCode": -3, 
+                    "content": [
+                        {
+                            "buffer": None
+                        },
+                    ], 
+                    "info": "input content is not valid", 
+                }
+
+                print(f"errCode: -3, input content is not valid!")
+                return ret_dict
+
+            # Preprocess input batch
+            torch.cuda.synchronize()
+
+            a = datetime.now()
+            
+            try:
+                model_kwargs_tmp = data_preprocess_server(
+                                        args, image_path, driving_audio_path, prompt, feature_extractor
+                                        )
+            except:
+                ret_dict = {
+                    "errCode": -2,         
+                    "content": [
+                            {
+                                "buffer": None
+                            },
+                        ],
+                    "info": "failed to preprocess input data"
+                }
+                print(f"errCode: -2, preprocess failed!")
+                return ret_dict
+
+            text_prompt = model_kwargs_tmp["text_prompt"]
+            audio_path = model_kwargs_tmp["audio_path"]
+            image_path = model_kwargs_tmp["image_path"]
+            fps = model_kwargs_tmp["fps"]
+            audio_prompts = model_kwargs_tmp["audio_prompts"]
+            audio_len = model_kwargs_tmp["audio_len"]
+            motion_bucket_id_exps = model_kwargs_tmp["motion_bucket_id_exps"]
+            motion_bucket_id_heads = model_kwargs_tmp["motion_bucket_id_heads"]
+            pixel_value_ref = model_kwargs_tmp["pixel_value_ref"]
+            pixel_value_ref_llava = model_kwargs_tmp["pixel_value_ref_llava"]
+            
+
+
+            torch.cuda.synchronize()
+            b = datetime.now()
+            preprocess_time = (b - a).total_seconds()
+            print("="*100)
+            print("preprocess time :", preprocess_time)
+            print("="*100)
+            
+        else:
+            text_prompt = None
+            audio_path = None
+            image_path = None
+            fps = None
+            audio_prompts = None
+            audio_len = None
+            motion_bucket_id_exps = None
+            motion_bucket_id_heads = None
+            pixel_value_ref = None
+            pixel_value_ref_llava = None
+
+    except:
+        traceback.print_exc()
+        if rank == 0:
+            ret_dict = {
+                "errCode": -1,         # Failed to generate video
+                "content":[
+                    {
+                        "buffer": None
+                    }
+                ],
+                "info": "failed to preprocess",
+            }
+            return ret_dict
+
+    try:
+        broadcast_params = [
+            text_prompt,
+            audio_path,
+            image_path,
+            fps,
+            audio_prompts,
+            audio_len,
+            motion_bucket_id_exps,
+            motion_bucket_id_heads,
+            pixel_value_ref,
+            pixel_value_ref_llava,
+        ]
+        dist.broadcast_object_list(broadcast_params, src=0)
+        outputs = generate_image_parallel(*broadcast_params)
+
+        if rank == 0:
+            samples = outputs["samples"]
+            sample = samples[0].unsqueeze(0)
+
+            sample = sample[:, :, :audio_len[0]]
+            
+            video = sample[0].permute(1, 2, 3, 0).clamp(0, 1).numpy()
+            video = (video * 255.).astype(np.uint8)
+
+            output_dict = {
+                "err_code": 0, 
+                "err_msg": "succeed", 
+                "video": video, 
+                "audio": input_dict.get("audio_path", None), 
+                "save_fps": save_fps, 
+            }
+
+            ret_dict = process_output_dict(output_dict)
+            return ret_dict
+    
+    except:
+        traceback.print_exc()
+        if rank == 0:
+            ret_dict = {
+                "errCode": -1,         # Failed to generate video
+                "content":[
+                    {
+                        "buffer": None
+                    }
+                ],
+                "info": "failed to generate video",
+            }
+            return ret_dict
+        
+    return None
+    
+def generate_image_parallel(text_prompt,
+                    audio_path,
+                    image_path,
+                    fps,
+                    audio_prompts,
+                    audio_len,
+                    motion_bucket_id_exps,
+                    motion_bucket_id_heads,
+                    pixel_value_ref,
+                    pixel_value_ref_llava
+                    ):
+    if nccl_info.sp_size > 1:
+        device = torch.device(f"cuda:{torch.distributed.get_rank()}")
+
+    batch = {
+        "text_prompt": text_prompt,
+        "audio_path": audio_path,
+        "image_path": image_path,
+        "fps": fps,
+        "audio_prompts": audio_prompts,
+        "audio_len": audio_len,
+        "motion_bucket_id_exps": motion_bucket_id_exps,
+        "motion_bucket_id_heads": motion_bucket_id_heads,
+        "pixel_value_ref": pixel_value_ref,
+        "pixel_value_ref_llava": pixel_value_ref_llava
+    }
+
+    samples = hunyuan_sampler.predict(args, batch, wav2vec, feature_extractor, align_instance)
+    return samples
+
+def worker_loop():
+    while True:
+        predict_wrap()
+        
+
+if __name__ == "__main__":
+    audio_args = parse_args()
+    initialize_distributed(audio_args.seed)
+    hunyuan_sampler = HunyuanVideoSampler.from_pretrained(
+        audio_args.ckpt, args=audio_args)
+    args = hunyuan_sampler.args
+    
+    rank = local_rank = 0
+    device = torch.device("cuda")
+    if nccl_info.sp_size > 1:
+        device = torch.device(f"cuda:{torch.distributed.get_rank()}")
+        rank = local_rank = torch.distributed.get_rank()
+
+    feature_extractor = AutoFeatureExtractor.from_pretrained(f"{MODEL_OUTPUT_PATH}/ckpts/whisper-tiny/")
+    wav2vec = WhisperModel.from_pretrained(f"{MODEL_OUTPUT_PATH}/ckpts/whisper-tiny/").to(device=device, dtype=torch.float32)
+    wav2vec.requires_grad_(False)
+
+
+    BASE_DIR = f'{MODEL_OUTPUT_PATH}/ckpts/det_align/'
+    det_path = os.path.join(BASE_DIR, 'detface.pt')    
+    align_instance = AlignImage("cuda", det_path=det_path)
+
+
+
+    if rank == 0:
+        uvicorn.run(app, host="0.0.0.0", port=80)
+    else:
+        worker_loop()
+    

hymm_gradio/gradio_audio.py

@@ -0,0 +1,122 @@
+import os
+import cv2
+import glob
+import json
+import datetime
+import requests
+import gradio as gr
+from tool_for_end2end import *
+
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+DATADIR = './temp'
+_HEADER_ = '''
+<div style="text-align: center; max-width: 650px; margin: 0 auto;">
+    <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">Tencent HunyuanVideo-Avatar Demo</h1>
+</div>
+
+''' 
+# flask url
+URL = "http://127.0.0.1:80/predict2"
+
+def post_and_get(audio_input, id_image, prompt):
+    now = datetime.datetime.now().isoformat()
+    imgdir = os.path.join(DATADIR, 'reference')
+    videodir = os.path.join(DATADIR, 'video')
+    imgfile = os.path.join(imgdir, now + '.png')
+    output_video_path = os.path.join(videodir, now + '.mp4')
+
+
+    os.makedirs(imgdir, exist_ok=True)
+    os.makedirs(videodir, exist_ok=True)
+    cv2.imwrite(imgfile, id_image[:,:,::-1])
+
+    proxies = {
+        "http": None, 
+        "https": None,
+    }
+    
+    files = {
+        "image_buffer": encode_image_to_base64(imgfile), 
+        "audio_buffer":  encode_wav_to_base64(audio_input),
+        "text": prompt,
+        "save_fps": 25, 
+    }
+    r = requests.get(URL, data = json.dumps(files), proxies=proxies)
+    ret_dict = json.loads(r.text)
+    print(ret_dict["info"])
+    save_video_base64_to_local(
+        video_path=None, 
+        base64_buffer=ret_dict["content"][0]["buffer"], 
+        output_video_path=output_video_path)
+
+
+    return output_video_path
+
+def create_demo():
+    
+    with gr.Blocks() as demo:
+        gr.Markdown(_HEADER_)
+        with gr.Tab('语音数字人驱动'):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Group():
+                        prompt = gr.Textbox(label="Prompt", value="a man is speaking.")
+  
+                    audio_input = gr.Audio(sources=["upload"],
+                                       type="filepath",
+                                       label="Upload Audio",
+                                       elem_classes="media-upload",
+                                       scale=1)
+                    id_image = gr.Image(label="Input reference image", height=480)
+
+                with gr.Column(scale=2):
+                    with gr.Group():
+                        output_image = gr.Video(label="Generated Video")
+                        
+
+                with gr.Column(scale=1):
+                    generate_btn = gr.Button("Generate")
+
+            generate_btn.click(fn=post_and_get,
+                inputs=[audio_input, id_image, prompt],
+                outputs=[output_image],
+            )
+            
+            # quick_prompts = [[x] for x in glob.glob('./assets/images/*.png')]
+            # example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Other object', samples_per_page=1000, components=[id_image])
+            # example_quick_prompts.click(lambda x: x[0], inputs=example_quick_prompts, outputs=id_image, show_progress=False, queue=False)
+            # with gr.Row(), gr.Column():
+            #     gr.Markdown("## Examples")
+            #     example_inps = [
+            #         [
+            #             'A woman is drinking coffee at a café.',
+            #             './assets/images/seg_woman_01.png',
+            #             1280, 720, 30, 129, 7.5, 13, 1024,
+            #             "assets/videos/seg_woman_01.mp4"
+            #         ],
+            #         [
+            #             'In a cubicle of an office building, a woman focuses intently on the computer screen, typing rapidly on the keyboard, surrounded by piles of documents.',
+            #             './assets/images/seg_woman_03.png',
+            #             1280, 720, 30, 129, 7.5, 13, 1025,
+            #             "./assets/videos/seg_woman_03.mp4"
+            #         ],
+            #         [
+            #             'A man walks across an ancient stone bridge holding an umbrella, raindrops tapping against it.',
+            #             './assets/images/seg_man_01.png',
+            #             1280, 720, 30, 129, 7.5, 13, 1025,
+            #             "./assets/videos/seg_man_01.mp4"
+            #         ],
+            #         [
+            #             'During a train journey, a man admires the changing scenery through the window.',
+            #             './assets/images/seg_man_02.png',
+            #             1280, 720, 30, 129, 7.5, 13, 1026,
+            #             "./assets/videos/seg_man_02.mp4"
+            #         ]
+            #     ]
+            #     gr.Examples(examples=example_inps, inputs=[prompt, id_image, width, height, num_steps, num_frames, guidance, flow_shift, seed, output_image],)
+    return demo
+
+if __name__ == "__main__":
+    allowed_paths = ['/']
+    demo = create_demo()
+    demo.launch(server_name='0.0.0.0', server_port=8080, share=True, allowed_paths=allowed_paths)

hymm_gradio/tool_for_end2end.py

@@ -0,0 +1,325 @@
+import os
+import io
+import math
+import uuid
+import base64
+import imageio
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+from copy import deepcopy
+from einops import rearrange
+import torchvision.transforms as transforms
+from torchvision.transforms import ToPILImage
+from hymm_sp.data_kits.audio_dataset import get_audio_feature
+from hymm_sp.data_kits.ffmpeg_utils import save_video
+
+TEMP_DIR = "./temp"
+if not os.path.exists(TEMP_DIR):
+    os.makedirs(TEMP_DIR, exist_ok=True)
+
+
+def data_preprocess_server(args, image_path, audio_path, prompts, feature_extractor):
+    llava_transform = transforms.Compose(
+            [
+                transforms.Resize((336, 336), interpolation=transforms.InterpolationMode.BILINEAR), 
+                transforms.ToTensor(), 
+                transforms.Normalize((0.48145466, 0.4578275, 0.4082107), (0.26862954, 0.26130258, 0.27577711)),
+            ]
+        )
+    
+    """ 生成prompt """
+    if prompts is None:
+        prompts = "Authentic, Realistic, Natural, High-quality, Lens-Fixed." 
+    else:
+        prompts = "Authentic, Realistic, Natural, High-quality, Lens-Fixed, " + prompts
+
+    fps = 25
+    
+    img_size = args.image_size
+    ref_image = Image.open(image_path).convert('RGB')
+    
+    # Resize reference image
+    w, h = ref_image.size
+    scale = img_size / min(w, h)
+    new_w = round(w * scale / 64) * 64
+    new_h = round(h * scale / 64) * 64
+
+    if img_size == 704:
+        img_size_long = 1216
+    if new_w * new_h > img_size * img_size_long:
+        scale = math.sqrt(img_size * img_size_long / w / h)
+        new_w = round(w * scale / 64) * 64
+        new_h = round(h * scale / 64) * 64
+
+    ref_image = ref_image.resize((new_w, new_h), Image.LANCZOS)
+    
+    ref_image = np.array(ref_image)
+    ref_image = torch.from_numpy(ref_image)
+        
+    audio_input, audio_len = get_audio_feature(feature_extractor, audio_path)
+    audio_prompts = audio_input[0]
+    
+    motion_bucket_id_heads = np.array([25] * 4)
+    motion_bucket_id_exps = np.array([30] * 4)
+    motion_bucket_id_heads = torch.from_numpy(motion_bucket_id_heads)
+    motion_bucket_id_exps = torch.from_numpy(motion_bucket_id_exps)
+    fps = torch.from_numpy(np.array(fps))
+    
+    to_pil = ToPILImage()
+    pixel_value_ref = rearrange(ref_image.clone().unsqueeze(0), "b h w c -> b c h w")   # (b c h w)
+    
+    pixel_value_ref_llava = [llava_transform(to_pil(image)) for image in pixel_value_ref]
+    pixel_value_ref_llava = torch.stack(pixel_value_ref_llava, dim=0)
+
+    batch = {
+        "text_prompt": [prompts],
+        "audio_path": [audio_path],
+        "image_path": [image_path],
+        "fps": fps.unsqueeze(0).to(dtype=torch.float16),
+        "audio_prompts": audio_prompts.unsqueeze(0).to(dtype=torch.float16),
+        "audio_len": [audio_len],
+        "motion_bucket_id_exps": motion_bucket_id_exps.unsqueeze(0),
+        "motion_bucket_id_heads": motion_bucket_id_heads.unsqueeze(0),
+        "pixel_value_ref": pixel_value_ref.unsqueeze(0).to(dtype=torch.float16),
+        "pixel_value_ref_llava": pixel_value_ref_llava.unsqueeze(0).to(dtype=torch.float16)
+    }
+
+    return batch
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8, quality=8):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    outputs = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = torch.clamp(x,0,1)
+        x = (x * 255).numpy().astype(np.uint8)
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    imageio.mimsave(path, outputs, fps=fps, quality=quality)
+
+def encode_image_to_base64(image_path):
+    try:
+        with open(image_path, 'rb') as image_file:
+            image_data = image_file.read()
+        encoded_data = base64.b64encode(image_data).decode('utf-8')
+        print(f"Image file '{image_path}' has been successfully encoded to Base64.")
+        return encoded_data
+    
+    except Exception as e:
+        print(f"Error encoding image: {e}")
+        return None
+
+def encode_video_to_base64(video_path):
+    try:
+        with open(video_path, 'rb') as video_file:
+            video_data = video_file.read()
+        encoded_data = base64.b64encode(video_data).decode('utf-8')
+        print(f"Video file '{video_path}' has been successfully encoded to Base64.")
+        return encoded_data
+    
+    except Exception as e:
+        print(f"Error encoding video: {e}")
+        return None
+    
+def encode_wav_to_base64(wav_path):
+    try:
+        with open(wav_path, 'rb') as audio_file:
+            audio_data = audio_file.read()
+        encoded_data = base64.b64encode(audio_data).decode('utf-8')
+        print(f"Audio file '{wav_path}' has been successfully encoded to Base64.")
+        return encoded_data
+    
+    except Exception as e:
+        print(f"Error encoding audio: {e}")
+        return None
+    
+def encode_pkl_to_base64(pkl_path):
+    try:
+        with open(pkl_path, 'rb') as pkl_file:
+            pkl_data = pkl_file.read()
+        
+        encoded_data = base64.b64encode(pkl_data).decode('utf-8')
+        
+        print(f"Pickle file '{pkl_path}' has been successfully encoded to Base64.")
+        return encoded_data
+
+    except Exception as e:
+        print(f"Error encoding pickle: {e}")
+        return None
+      
+def decode_base64_to_image(base64_buffer_str):
+    try:
+        image_data = base64.b64decode(base64_buffer_str)
+        image = Image.open(io.BytesIO(image_data))
+        image_array = np.array(image)
+        print(f"Image Base64 string has beed succesfully decoded to image.")
+        return image_array
+    except Exception as e:
+        print(f"Error encdecodingoding image: {e}")
+        return None
+    
+def decode_base64_to_video(base64_buffer_str):
+    try:
+        video_data = base64.b64decode(base64_buffer_str)
+        video_bytes = io.BytesIO(video_data)
+        video_bytes.seek(0)
+        video_reader = imageio.get_reader(video_bytes, 'ffmpeg')
+        video_frames = [frame for frame in video_reader]
+        return video_frames
+    except Exception as e:
+        print(f"Error decoding video: {e}")
+        return None
+
+    
+def save_video_base64_to_local(video_path=None, base64_buffer=None, output_video_path=None):
+    if video_path is not None and base64_buffer is None:
+        video_buffer_base64 = encode_video_to_base64(video_path)
+    elif video_path is None and base64_buffer is not None:
+        video_buffer_base64 = deepcopy(base64_buffer)
+    else:
+        print("Please pass either 'video_path' or 'base64_buffer'")
+        return None
+    
+    if video_buffer_base64 is not None:
+        video_data = base64.b64decode(video_buffer_base64)
+        if output_video_path is None:
+            uuid_string = str(uuid.uuid4())
+            temp_video_path = f'{TEMP_DIR}/{uuid_string}.mp4'
+        else:
+            temp_video_path = output_video_path
+        with open(temp_video_path, 'wb') as video_file:
+            video_file.write(video_data)
+        return temp_video_path
+    else:
+        return None
+    
+def save_audio_base64_to_local(audio_path=None, base64_buffer=None):
+    if audio_path is not None and base64_buffer is None:
+        audio_buffer_base64 = encode_wav_to_base64(audio_path)
+    elif audio_path is None and base64_buffer is not None:
+        audio_buffer_base64 = deepcopy(base64_buffer)
+    else:
+        print("Please pass either 'audio_path' or 'base64_buffer'")
+        return None
+    
+    if audio_buffer_base64 is not None:
+        audio_data = base64.b64decode(audio_buffer_base64)
+        uuid_string = str(uuid.uuid4())
+        temp_audio_path = f'{TEMP_DIR}/{uuid_string}.wav'
+        with open(temp_audio_path, 'wb') as audio_file:
+            audio_file.write(audio_data)
+        return temp_audio_path
+    else:
+        return None
+    
+def save_pkl_base64_to_local(pkl_path=None, base64_buffer=None):
+    if pkl_path is not None and base64_buffer is None:
+        pkl_buffer_base64 = encode_pkl_to_base64(pkl_path)
+    elif pkl_path is None and base64_buffer is not None:
+        pkl_buffer_base64 = deepcopy(base64_buffer)
+    else:
+        print("Please pass either 'pkl_path' or 'base64_buffer'")
+        return None
+    
+    if pkl_buffer_base64 is not None:
+        pkl_data = base64.b64decode(pkl_buffer_base64)
+        uuid_string = str(uuid.uuid4())
+        temp_pkl_path = f'{TEMP_DIR}/{uuid_string}.pkl'
+        with open(temp_pkl_path, 'wb') as pkl_file:
+            pkl_file.write(pkl_data)
+        return temp_pkl_path
+    else:
+        return None
+    
+def remove_temp_fles(input_dict):
+    for key, val in input_dict.items():
+        if "_path" in key and val is not None and os.path.exists(val):
+            os.remove(val)
+            print(f"Remove temporary {key} from {val}")
+
+def process_output_dict(output_dict):
+
+    uuid_string = str(uuid.uuid4())
+    temp_video_path = f'{TEMP_DIR}/{uuid_string}.mp4'
+    save_video(output_dict["video"], temp_video_path, fps=output_dict.get("save_fps", 25))
+
+    # Add audio
+    if output_dict["audio"] is not None and os.path.exists(output_dict["audio"]):
+        output_path = temp_video_path
+        audio_path = output_dict["audio"]
+        save_path = temp_video_path.replace(".mp4", "_audio.mp4")
+        print('='*100)
+        print(f"output_path = {output_path}\n audio_path = {audio_path}\n save_path = {save_path}")
+        os.system(f"ffmpeg -i '{output_path}' -i '{audio_path}' -shortest '{save_path}' -y -loglevel quiet; rm '{output_path}'")
+    else:
+        save_path = temp_video_path
+
+    video_base64_buffer = encode_video_to_base64(save_path)
+
+    encoded_output_dict = {
+        "errCode": output_dict["err_code"], 
+        "content": [
+                    {
+                        "buffer": video_base64_buffer
+                    },
+                ],
+        "info":output_dict["err_msg"],
+    }
+    
+    
+
+    return encoded_output_dict
+
+
+def save_image_base64_to_local(image_path=None, base64_buffer=None):
+    # Encode image to base64 buffer
+    if image_path is not None and base64_buffer is None:
+        image_buffer_base64 = encode_image_to_base64(image_path)
+    elif image_path is None and base64_buffer is not None:
+        image_buffer_base64 = deepcopy(base64_buffer)
+    else:
+        print("Please pass either 'image_path' or 'base64_buffer'")
+        return None
+    
+    # Decode base64 buffer and save to local disk
+    if image_buffer_base64 is not None:
+        image_data = base64.b64decode(image_buffer_base64)
+        uuid_string = str(uuid.uuid4())
+        temp_image_path = f'{TEMP_DIR}/{uuid_string}.png'
+        with open(temp_image_path, 'wb') as image_file:
+            image_file.write(image_data)
+        return temp_image_path
+    else:
+        return None
+    
+def process_input_dict(input_dict):
+    
+    decoded_input_dict = {}
+   
+    decoded_input_dict["save_fps"] = input_dict.get("save_fps", 25)
+
+    image_base64_buffer = input_dict.get("image_buffer", None)
+    if image_base64_buffer is not None:
+        decoded_input_dict["image_path"] = save_image_base64_to_local(
+            image_path=None, 
+            base64_buffer=image_base64_buffer)
+    else:
+        decoded_input_dict["image_path"] = None
+    
+    audio_base64_buffer = input_dict.get("audio_buffer", None)
+    if audio_base64_buffer is not None:
+        decoded_input_dict["audio_path"] = save_audio_base64_to_local(
+            audio_path=None, 
+            base64_buffer=audio_base64_buffer)
+    else:
+        decoded_input_dict["audio_path"] = None
+    
+    decoded_input_dict["prompt"] = input_dict.get("text", None)
+        
+    return decoded_input_dict

hymm_sp/config.py

@@ -0,0 +1,142 @@
+import argparse
+from hymm_sp.constants import *
+import re
+import collections.abc
+
+def as_tuple(x):
+    if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+        return tuple(x)
+    if x is None or isinstance(x, (int, float, str)):
+        return (x,)
+    else:
+        raise ValueError(f"Unknown type {type(x)}")
+
+def parse_args(namespace=None):
+    parser = argparse.ArgumentParser(description="Hunyuan Multimodal training/inference script")
+    parser = add_extra_args(parser)
+    args = parser.parse_args(namespace=namespace)
+    args = sanity_check_args(args)
+    return args
+
+def add_extra_args(parser: argparse.ArgumentParser):
+    parser = add_network_args(parser)
+    parser = add_extra_models_args(parser)
+    parser = add_denoise_schedule_args(parser)
+    parser = add_evaluation_args(parser)
+    return parser
+
+def add_network_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Network")
+    group.add_argument("--model", type=str, default="HYVideo-T/2",
+                       help="Model architecture to use. It it also used to determine the experiment directory.")
+    group.add_argument("--latent-channels", type=str, default=None,
+                       help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, "
+                            "it still needs to match the latent channels of the VAE model.")
+    group.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.")
+    return parser
+
+def add_extra_models_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Extra Models (VAE, Text Encoder, Tokenizer)")
+
+    # VAE
+    group.add_argument("--vae", type=str, default="884-16c-hy0801",  help="Name of the VAE model.")
+    group.add_argument("--vae-precision", type=str, default="fp16", 
+                       help="Precision mode for the VAE model.")
+    group.add_argument("--vae-tiling", action="store_true", default=True, help="Enable tiling for the VAE model.")
+    group.add_argument("--text-encoder", type=str, default="llava-llama-3-8b", choices=list(TEXT_ENCODER_PATH),
+                       help="Name of the text encoder model.")
+    group.add_argument("--text-encoder-precision", type=str, default="fp16", choices=PRECISIONS,
+                       help="Precision mode for the text encoder model.")
+    group.add_argument("--text-states-dim", type=int, default=4096, help="Dimension of the text encoder hidden states.")
+    group.add_argument("--text-len", type=int, default=256, help="Maximum length of the text input.")
+    group.add_argument("--tokenizer", type=str, default="llava-llama-3-8b", choices=list(TOKENIZER_PATH),
+                       help="Name of the tokenizer model.")
+    group.add_argument("--text-encoder-infer-mode", type=str, default="encoder", choices=["encoder", "decoder"],
+                       help="Inference mode for the text encoder model. It should match the text encoder type. T5 and "
+                            "CLIP can only work in 'encoder' mode, while Llava/GLM can work in both modes.")
+    group.add_argument("--prompt-template-video", type=str, default='li-dit-encode-video', choices=PROMPT_TEMPLATE,
+                       help="Video prompt template for the decoder-only text encoder model.")
+    group.add_argument("--hidden-state-skip-layer", type=int, default=2,
+                       help="Skip layer for hidden states.")
+    group.add_argument("--apply-final-norm", action="store_true",
+                       help="Apply final normalization to the used text encoder hidden states.")
+
+    # - CLIP
+    group.add_argument("--text-encoder-2", type=str, default='clipL', choices=list(TEXT_ENCODER_PATH),
+                       help="Name of the second text encoder model.")
+    group.add_argument("--text-encoder-precision-2", type=str, default="fp16", choices=PRECISIONS,
+                       help="Precision mode for the second text encoder model.")
+    group.add_argument("--text-states-dim-2", type=int, default=768,
+                       help="Dimension of the second text encoder hidden states.")
+    group.add_argument("--tokenizer-2", type=str, default='clipL', choices=list(TOKENIZER_PATH),
+                       help="Name of the second tokenizer model.")
+    group.add_argument("--text-len-2", type=int, default=77, help="Maximum length of the second text input.")
+    group.set_defaults(use_attention_mask=True)
+    group.add_argument("--text-projection", type=str, default="single_refiner", choices=TEXT_PROJECTION,
+                       help="A projection layer for bridging the text encoder hidden states and the diffusion model "
+                            "conditions.")
+    return parser
+
+
+def add_denoise_schedule_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Denoise schedule")
+    group.add_argument("--flow-shift-eval-video", type=float, default=None, help="Shift factor for flow matching schedulers when using video data.")
+    group.add_argument("--flow-reverse", action="store_true", default=True, help="If reverse, learning/sampling from t=1 -> t=0.")
+    group.add_argument("--flow-solver", type=str, default="euler", help="Solver for flow matching.")
+    group.add_argument("--use-linear-quadratic-schedule", action="store_true", help="Use linear quadratic schedule for flow matching."
+                                                    "Follow MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)")
+    group.add_argument("--linear-schedule-end", type=int, default=25, help="End step for linear quadratic schedule for flow matching.")
+    return parser
+
+def add_evaluation_args(parser: argparse.ArgumentParser):
+    group = parser.add_argument_group(title="Validation Loss Evaluation")
+    parser.add_argument("--precision", type=str, default="bf16", choices=PRECISIONS,
+                    help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.")
+    parser.add_argument("--reproduce", action="store_true",
+                       help="Enable reproducibility by setting random seeds and deterministic algorithms.")
+    parser.add_argument("--ckpt", type=str, help="Path to the checkpoint to evaluate.")
+    parser.add_argument("--load-key", type=str, default="module", choices=["module", "ema"],
+                       help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.")
+    parser.add_argument("--cpu-offload", action="store_true", help="Use CPU offload for the model load.")
+    parser.add_argument("--infer-min", action="store_true", help="infer 5s.")
+    group.add_argument( "--use-fp8", action="store_true", help="Enable use fp8 for inference acceleration.")
+    group.add_argument("--video-size", type=int, nargs='+', default=512,
+                        help="Video size for training. If a single value is provided, it will be used for both width "
+                            "and height. If two values are provided, they will be used for width and height "
+                            "respectively.")
+    group.add_argument("--sample-n-frames", type=int, default=1,
+                       help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1")
+    group.add_argument("--infer-steps", type=int, default=100, help="Number of denoising steps for inference.")
+    group.add_argument("--val-disable-autocast", action="store_true",
+                       help="Disable autocast for denoising loop and vae decoding in pipeline sampling.")
+    group.add_argument("--num-images", type=int, default=1, help="Number of images to generate for each prompt.")
+    group.add_argument("--seed", type=int, default=1024, help="Seed for evaluation.")
+    group.add_argument("--save-path-suffix", type=str, default="", help="Suffix for the directory of saved samples.")
+    group.add_argument("--pos-prompt", type=str, default='', help="Prompt for sampling during evaluation.")
+    group.add_argument("--neg-prompt", type=str, default='', help="Negative prompt for sampling during evaluation.")
+    group.add_argument("--image-size", type=int, default=704)
+    group.add_argument("--pad-face-size", type=float, default=0.7, help="Pad bbox for face align.")
+    group.add_argument("--image-path", type=str, default="",  help="")
+    group.add_argument("--save-path", type=str, default=None, help="Path to save the generated samples.")
+    group.add_argument("--input", type=str, default=None, help="test data.")
+    group.add_argument("--item-name", type=str, default=None, help="")
+    group.add_argument("--cfg-scale", type=float, default=7.5, help="Classifier free guidance scale.")
+    group.add_argument("--ip-cfg-scale", type=float, default=0, help="Classifier free guidance scale.")
+    group.add_argument("--use-deepcache", type=int, default=1)
+    return parser
+
+def sanity_check_args(args):
+    # VAE channels
+    vae_pattern = r"\d{2,3}-\d{1,2}c-\w+"
+    if not re.match(vae_pattern, args.vae):
+        raise ValueError(
+            f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'."
+        )
+    vae_channels = int(args.vae.split("-")[1][:-1])
+    if args.latent_channels is None:
+        args.latent_channels = vae_channels
+    if vae_channels != args.latent_channels:
+        raise ValueError(
+            f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels})."
+        )
+    return args

hymm_sp/constants.py

@@ -0,0 +1,59 @@
+import os
+import torch
+
+__all__ = [
+    "PROMPT_TEMPLATE", "MODEL_BASE", "PRECISION_TO_TYPE",
+    "PRECISIONS", "VAE_PATH", "TEXT_ENCODER_PATH", "TOKENIZER_PATH",
+    "TEXT_PROJECTION",
+]
+
+# =================== Constant Values =====================
+
+PRECISION_TO_TYPE = {
+    'fp32': torch.float32,
+    'fp16': torch.float16,
+    'bf16': torch.bfloat16,
+}
+
+PROMPT_TEMPLATE_ENCODE_VIDEO = (
+    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
+    "1. The main content and theme of the video."
+    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
+    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
+    "4. background environment, light, style and atmosphere."
+    "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
+    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
+)
+
+PROMPT_TEMPLATE = {
+    "li-dit-encode-video": {"template": PROMPT_TEMPLATE_ENCODE_VIDEO, "crop_start": 95},
+}
+
+# ======================= Model ======================
+PRECISIONS = {"fp32", "fp16", "bf16"}
+
+# =================== Model Path =====================
+MODEL_BASE = os.getenv("MODEL_BASE")
+MODEL_BASE=f"{MODEL_BASE}/ckpts"
+
+# 3D VAE
+VAE_PATH = {
+    "884-16c-hy0801": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae",
+}
+
+# Text Encoder
+TEXT_ENCODER_PATH = {
+    "clipL": f"{MODEL_BASE}/text_encoder_2",
+    "llava-llama-3-8b": f"{MODEL_BASE}/llava_llama_image",
+}
+
+# Tokenizer
+TOKENIZER_PATH = {
+    "clipL": f"{MODEL_BASE}/text_encoder_2",
+    "llava-llama-3-8b":f"{MODEL_BASE}/llava_llama_image",
+}
+
+TEXT_PROJECTION = {
+    "linear",                               # Default, an nn.Linear() layer
+    "single_refiner",                       # Single TokenRefiner. Refer to LI-DiT
+}

hymm_sp/data_kits/audio_dataset.py

@@ -0,0 +1,170 @@
+import os
+import cv2
+import math
+import json
+import torch
+import random
+import librosa
+import traceback
+import torchvision
+import numpy as np
+import pandas as pd
+from PIL import Image
+from einops import rearrange
+from torch.utils.data import Dataset
+from decord import VideoReader, cpu
+from transformers import CLIPImageProcessor
+import torchvision.transforms as transforms
+from torchvision.transforms import ToPILImage
+
+
+
+def get_audio_feature(feature_extractor, audio_path):
+    audio_input, sampling_rate = librosa.load(audio_path, sr=16000)
+    assert sampling_rate == 16000
+
+    audio_features = []
+    window = 750*640
+    for i in range(0, len(audio_input), window):
+        audio_feature = feature_extractor(audio_input[i:i+window], 
+                                        sampling_rate=sampling_rate, 
+                                        return_tensors="pt", 
+                                        ).input_features
+        audio_features.append(audio_feature)
+
+    audio_features = torch.cat(audio_features, dim=-1)
+    return audio_features, len(audio_input) // 640
+
+
+class VideoAudioTextLoaderVal(Dataset):
+    def __init__(
+        self, 
+        image_size: int,
+        meta_file: str, 
+        **kwargs,
+    ):
+        super().__init__()
+        self.meta_file = meta_file
+        self.image_size = image_size
+        self.text_encoder = kwargs.get("text_encoder", None)            # llava_text_encoder
+        self.text_encoder_2 = kwargs.get("text_encoder_2", None)        # clipL_text_encoder
+        self.feature_extractor = kwargs.get("feature_extractor", None)
+        self.meta_files = []
+    
+        csv_data = pd.read_csv(meta_file)
+        for idx in range(len(csv_data)):
+            self.meta_files.append(
+                {
+                    "videoid": str(csv_data["videoid"][idx]),
+                    "image_path": str(csv_data["image"][idx]), 
+                    "audio_path": str(csv_data["audio"][idx]), 
+                    "prompt": str(csv_data["prompt"][idx]), 
+                    "fps": float(csv_data["fps"][idx])
+                }
+            )
+        
+        self.llava_transform = transforms.Compose(
+            [
+                transforms.Resize((336, 336), interpolation=transforms.InterpolationMode.BILINEAR), 
+                transforms.ToTensor(), 
+                transforms.Normalize((0.48145466, 0.4578275, 0.4082107), (0.26862954, 0.26130258, 0.27577711)),
+            ]
+        )
+        self.clip_image_processor = CLIPImageProcessor()
+        
+        self.device = torch.device("cuda")
+        self.weight_dtype = torch.float16
+
+        
+    def __len__(self):
+        return len(self.meta_files)
+
+    @staticmethod
+    def get_text_tokens(text_encoder, description, dtype_encode="video"):
+        text_inputs = text_encoder.text2tokens(description, data_type=dtype_encode)
+        text_ids = text_inputs["input_ids"].squeeze(0)
+        text_mask = text_inputs["attention_mask"].squeeze(0)
+        return text_ids, text_mask
+    
+    def get_batch_data(self, idx):
+        meta_file = self.meta_files[idx]
+        videoid = meta_file["videoid"]
+        image_path = meta_file["image_path"]
+        audio_path = meta_file["audio_path"]
+        prompt = "Authentic, Realistic, Natural, High-quality, Lens-Fixed, " + meta_file["prompt"]
+        fps = meta_file["fps"]
+        
+        img_size = self.image_size
+        ref_image = Image.open(image_path).convert('RGB')
+        
+        # Resize reference image
+        w, h = ref_image.size
+        scale = img_size / min(w, h)
+        new_w = round(w * scale / 64) * 64
+        new_h = round(h * scale / 64) * 64
+
+        if img_size == 704:
+            img_size_long = 1216
+        if new_w * new_h > img_size * img_size_long:
+            import math
+            scale = math.sqrt(img_size * img_size_long / w / h)
+            new_w = round(w * scale / 64) * 64
+            new_h = round(h * scale / 64) * 64
+
+        ref_image = ref_image.resize((new_w, new_h), Image.LANCZOS)
+        
+        ref_image = np.array(ref_image)
+        ref_image = torch.from_numpy(ref_image)
+         
+        audio_input, audio_len = get_audio_feature(self.feature_extractor, audio_path)
+        audio_prompts = audio_input[0]
+        
+        motion_bucket_id_heads = np.array([25] * 4)
+        motion_bucket_id_exps = np.array([30] * 4)
+        motion_bucket_id_heads = torch.from_numpy(motion_bucket_id_heads)
+        motion_bucket_id_exps = torch.from_numpy(motion_bucket_id_exps)
+        fps = torch.from_numpy(np.array(fps))
+        
+        to_pil = ToPILImage()
+        pixel_value_ref = rearrange(ref_image.clone().unsqueeze(0), "b h w c -> b c h w")   # (b c h w)
+        
+        pixel_value_ref_llava = [self.llava_transform(to_pil(image)) for image in pixel_value_ref]
+        pixel_value_ref_llava = torch.stack(pixel_value_ref_llava, dim=0)
+        pixel_value_ref_clip = self.clip_image_processor(
+            images=Image.fromarray((pixel_value_ref[0].permute(1,2,0)).data.cpu().numpy().astype(np.uint8)), 
+            return_tensors="pt"
+        ).pixel_values[0]
+        pixel_value_ref_clip = pixel_value_ref_clip.unsqueeze(0)
+        
+        # Encode text prompts
+   
+        text_ids, text_mask = self.get_text_tokens(self.text_encoder, prompt)
+        text_ids_2, text_mask_2 = self.get_text_tokens(self.text_encoder_2, prompt)
+        
+        # Output batch
+        batch = {
+            "text_prompt": prompt,                         # 
+            "videoid": videoid, 
+            "pixel_value_ref": pixel_value_ref.to(dtype=torch.float16),                 # 参考图，用于vae提特征 (1, 3, h, w), 取值范围(0, 255)
+            "pixel_value_ref_llava": pixel_value_ref_llava.to(dtype=torch.float16),     # 参考图，用于llava提特征 (1, 3, 336, 336), 取值范围 = CLIP取值范围
+            "pixel_value_ref_clip": pixel_value_ref_clip.to(dtype=torch.float16),       # 参考图，用于clip_image_encoder提特征 (1, 3, 244, 244), 取值范围 = CLIP取值范围
+            "audio_prompts": audio_prompts.to(dtype=torch.float16), 
+            "motion_bucket_id_heads": motion_bucket_id_heads.to(dtype=text_ids.dtype), 
+            "motion_bucket_id_exps": motion_bucket_id_exps.to(dtype=text_ids.dtype), 
+            "fps": fps.to(dtype=torch.float16), 
+            "text_ids": text_ids.clone(),                                               # 对应llava_text_encoder
+            "text_mask": text_mask.clone(),                                             # 对应llava_text_encoder
+            "text_ids_2": text_ids_2.clone(),                                           # 对应clip_text_encoder
+            "text_mask_2": text_mask_2.clone(),                                         # 对应clip_text_encoder
+            "audio_len": audio_len,
+            "image_path": image_path, 
+            "audio_path": audio_path, 
+        }
+        return batch
+    
+    def __getitem__(self, idx):
+        return self.get_batch_data(idx)
+    
+
+        
+        

hymm_sp/data_kits/audio_preprocessor.py

@@ -0,0 +1,72 @@
+
+import os
+import cv2
+import json
+import time
+import decord
+import einops
+import librosa
+import torch
+import random
+import argparse
+import traceback
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+from einops import rearrange
+
+
+
+def get_facemask(ref_image, align_instance, area=1.25):
+    # ref_image: (b f c h w)
+    bsz, f, c, h, w = ref_image.shape
+    images = rearrange(ref_image, "b f c h w -> (b f) h w c").data.cpu().numpy().astype(np.uint8)
+    face_masks = []
+    for image in images:
+        image_pil = Image.fromarray(image).convert("RGB")
+        _, _, bboxes_list = align_instance(np.array(image_pil)[:,:,[2,1,0]], maxface=True)
+        try:
+            bboxSrc = bboxes_list[0]
+        except:
+            bboxSrc = [0, 0, w, h]
+        x1, y1, ww, hh = bboxSrc
+        x2, y2 = x1 + ww, y1 + hh
+        ww, hh = (x2-x1) * area, (y2-y1) * area
+        center = [(x2+x1)//2, (y2+y1)//2]
+        x1 = max(center[0] - ww//2, 0)
+        y1 = max(center[1] - hh//2, 0)
+        x2 = min(center[0] + ww//2, w)
+        y2 = min(center[1] + hh//2, h)
+        
+        face_mask = np.zeros_like(np.array(image_pil))
+        face_mask[int(y1):int(y2), int(x1):int(x2)] = 1.0
+        face_masks.append(torch.from_numpy(face_mask[...,:1]))
+    face_masks = torch.stack(face_masks, dim=0)     # (b*f, h, w, c)
+    face_masks = rearrange(face_masks, "(b f) h w c -> b c f h w", b=bsz, f=f)
+    face_masks = face_masks.to(device=ref_image.device, dtype=ref_image.dtype)
+    return face_masks
+
+
+def encode_audio(wav2vec, audio_feats, fps, num_frames=129):
+    if fps == 25:
+        start_ts = [0]
+        step_ts = [1]
+    elif fps == 12.5:
+        start_ts = [0]
+        step_ts = [2]
+    num_frames = min(num_frames, 400)
+    audio_feats = wav2vec.encoder(audio_feats.unsqueeze(0)[:, :, :3000], output_hidden_states=True).hidden_states
+    audio_feats = torch.stack(audio_feats, dim=2)
+    audio_feats = torch.cat([torch.zeros_like(audio_feats[:,:4]), audio_feats], 1)
+    
+    audio_prompts = []
+    for bb in range(1):
+        audio_feats_list = []
+        for f in range(num_frames):
+            cur_t = (start_ts[bb] + f * step_ts[bb]) * 2
+            audio_clip = audio_feats[bb:bb+1, cur_t: cur_t+10]
+            audio_feats_list.append(audio_clip)
+        audio_feats_list = torch.stack(audio_feats_list, 1)
+        audio_prompts.append(audio_feats_list)
+    audio_prompts = torch.cat(audio_prompts)
+    return audio_prompts

hymm_sp/data_kits/data_tools.py

@@ -0,0 +1,41 @@
+import os
+import cv2
+import torch
+import numpy as np
+import imageio
+import torchvision
+from einops import rearrange
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8, quality=8):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    outputs = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = torch.clamp(x,0,1)
+        x = (x * 255).numpy().astype(np.uint8)
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    imageio.mimsave(path, outputs, fps=fps, quality=quality)
+
+def pad_image(crop_img, size, color=(255, 255, 255), resize_ratio=1):
+    crop_h, crop_w = crop_img.shape[:2]
+    target_w, target_h = size
+    scale_h, scale_w = target_h / crop_h, target_w / crop_w
+    if scale_w > scale_h:
+        resize_h = int(target_h*resize_ratio)
+        resize_w = int(crop_w / crop_h * resize_h)
+    else:
+        resize_w = int(target_w*resize_ratio)
+        resize_h = int(crop_h / crop_w * resize_w)
+    crop_img = cv2.resize(crop_img, (resize_w, resize_h))
+    pad_left = (target_w - resize_w) // 2
+    pad_top = (target_h - resize_h) // 2
+    pad_right = target_w - resize_w - pad_left
+    pad_bottom = target_h - resize_h - pad_top
+    crop_img = cv2.copyMakeBorder(crop_img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=color)
+    return crop_img

hymm_sp/data_kits/face_align/__init__.py

@@ -0,0 +1 @@
+from .align import AlignImage

hymm_sp/data_kits/face_align/align.py

@@ -0,0 +1,34 @@
+import os
+import sys
+import torch
+from .detface import DetFace
+
+class AlignImage(object):
+    def __init__(self, device='cuda', det_path=''):
+        self.facedet = DetFace(pt_path=det_path, confThreshold=0.5, nmsThreshold=0.45, device=device)
+
+    @torch.no_grad()
+    def __call__(self, im, maxface=False):
+        bboxes, kpss, scores = self.facedet.detect(im)
+        face_num = bboxes.shape[0]
+
+        five_pts_list = []
+        scores_list = []
+        bboxes_list = []
+        for i in range(face_num):
+            five_pts_list.append(kpss[i].reshape(5,2))
+            scores_list.append(scores[i])
+            bboxes_list.append(bboxes[i])
+
+        if maxface and face_num>1:
+            max_idx = 0
+            max_area = (bboxes[0, 2])*(bboxes[0, 3])
+            for i in range(1, face_num):
+                area = (bboxes[i,2])*(bboxes[i,3])
+                if area>max_area:
+                    max_idx = i
+            five_pts_list = [five_pts_list[max_idx]]
+            scores_list = [scores_list[max_idx]]
+            bboxes_list = [bboxes_list[max_idx]]
+
+        return five_pts_list, scores_list, bboxes_list

hymm_sp/data_kits/face_align/detface.py

@@ -0,0 +1,283 @@
+# -*- coding: UTF-8 -*-
+import os
+import cv2
+import numpy as np
+import torch
+import torchvision
+
+
+def xyxy2xywh(x):
+    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
+    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
+    y[:, 2] = x[:, 2] - x[:, 0]  # width
+    y[:, 3] = x[:, 3] - x[:, 1]  # height
+    return y
+
+
+def xywh2xyxy(x):
+    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+    return y
+
+
+def box_iou(box1, box2):
+    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+    Arguments:
+        box1 (Tensor[N, 4])
+        box2 (Tensor[M, 4])
+    Returns:
+        iou (Tensor[N, M]): the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+    """
+
+    def box_area(box):
+        # box = 4xn
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area1 = box_area(box1.T)
+    area2 = box_area(box2.T)
+
+    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
+    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) -
+             torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
+    # iou = inter / (area1 + area2 - inter)
+    return inter / (area1[:, None] + area2 - inter)
+
+
+def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
+    # Rescale coords (xyxy) from img1_shape to img0_shape
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    coords[:, [0, 2]] -= pad[0]  # x padding
+    coords[:, [1, 3]] -= pad[1]  # y padding
+    coords[:, :4] /= gain
+    clip_coords(coords, img0_shape)
+    return coords
+
+
+def clip_coords(boxes, img_shape):
+    # Clip bounding xyxy bounding boxes to image shape (height, width)
+    boxes[:, 0].clamp_(0, img_shape[1])  # x1
+    boxes[:, 1].clamp_(0, img_shape[0])  # y1
+    boxes[:, 2].clamp_(0, img_shape[1])  # x2
+    boxes[:, 3].clamp_(0, img_shape[0])  # y2
+
+
+def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
+    # Rescale coords (xyxy) from img1_shape to img0_shape
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    coords[:, [0, 2, 4, 6, 8]] -= pad[0]  # x padding
+    coords[:, [1, 3, 5, 7, 9]] -= pad[1]  # y padding
+    coords[:, :10] /= gain
+    #clip_coords(coords, img0_shape)
+    coords[:, 0].clamp_(0, img0_shape[1])  # x1
+    coords[:, 1].clamp_(0, img0_shape[0])  # y1
+    coords[:, 2].clamp_(0, img0_shape[1])  # x2
+    coords[:, 3].clamp_(0, img0_shape[0])  # y2
+    coords[:, 4].clamp_(0, img0_shape[1])  # x3
+    coords[:, 5].clamp_(0, img0_shape[0])  # y3
+    coords[:, 6].clamp_(0, img0_shape[1])  # x4
+    coords[:, 7].clamp_(0, img0_shape[0])  # y4
+    coords[:, 8].clamp_(0, img0_shape[1])  # x5
+    coords[:, 9].clamp_(0, img0_shape[0])  # y5
+    return coords
+
+
+def show_results(img, xywh, conf, landmarks, class_num):
+    h,w,c = img.shape
+    tl = 1 or round(0.002 * (h + w) / 2) + 1  # line/font thickness
+    x1 = int(xywh[0] * w - 0.5 * xywh[2] * w)
+    y1 = int(xywh[1] * h - 0.5 * xywh[3] * h)
+    x2 = int(xywh[0] * w + 0.5 * xywh[2] * w)
+    y2 = int(xywh[1] * h + 0.5 * xywh[3] * h)
+    cv2.rectangle(img, (x1,y1), (x2, y2), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
+
+    clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
+
+    for i in range(5):
+        point_x = int(landmarks[2 * i] * w)
+        point_y = int(landmarks[2 * i + 1] * h)
+        cv2.circle(img, (point_x, point_y), tl+1, clors[i], -1)
+
+    tf = max(tl - 1, 1)  # font thickness
+    label = str(conf)[:5]
+    cv2.putText(img, label, (x1, y1 - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    return img
+
+
+def make_divisible(x, divisor):
+    # Returns x evenly divisible by divisor
+    return (x // divisor) * divisor
+
+
+def non_max_suppression_face(prediction, conf_thres=0.5, iou_thres=0.45, classes=None, agnostic=False, labels=()):
+    """Performs Non-Maximum Suppression (NMS) on inference results
+    Returns:
+         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
+    """
+
+    nc = prediction.shape[2] - 15  # number of classes
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
+    # time_limit = 10.0  # seconds to quit after
+    redundant = True  # require redundant detections
+    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)
+    merge = False  # use merge-NMS
+
+    # t = time.time()
+    output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Apply constraints
+        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
+        x = x[xc[xi]]  # confidence
+
+        # Cat apriori labels if autolabelling
+        if labels and len(labels[xi]):
+            l = labels[xi]
+            v = torch.zeros((len(l), nc + 15), device=x.device)
+            v[:, :4] = l[:, 1:5]  # box
+            v[:, 4] = 1.0  # conf
+            v[range(len(l)), l[:, 0].long() + 15] = 1.0  # cls
+            x = torch.cat((x, v), 0)
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Compute conf
+        x[:, 15:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+
+        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
+        box = xywh2xyxy(x[:, :4])
+
+        # Detections matrix nx6 (xyxy, conf, landmarks, cls)
+        if multi_label:
+            i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T
+            x = torch.cat((box[i], x[i, j + 15, None], x[i, 5:15] ,j[:, None].float()), 1)
+        else:  # best class only
+            conf, j = x[:, 15:].max(1, keepdim=True)
+            x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]
+
+        # Filter by class
+        if classes is not None:
+            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
+
+        # If none remain process next image
+        n = x.shape[0]  # number of boxes
+        if not n:
+            continue
+
+        # Batched NMS
+        c = x[:, 15:16] * (0 if agnostic else max_wh)  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+        #if i.shape[0] > max_det:  # limit detections
+        #    i = i[:max_det]
+        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
+            weights = iou * scores[None]  # box weights
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            if redundant:
+                i = i[iou.sum(1) > 1]  # require redundancy
+
+        output[xi] = x[i]
+        # if (time.time() - t) > time_limit:
+        #     break  # time limit exceeded
+
+    return output
+
+
+class DetFace():
+    def __init__(self, pt_path, confThreshold=0.5, nmsThreshold=0.45, device='cuda'):
+        assert os.path.exists(pt_path)
+
+        self.inpSize = 416
+        self.conf_thres = confThreshold
+        self.iou_thres = nmsThreshold
+        self.test_device = torch.device(device if torch.cuda.is_available() else "cpu")
+        self.model = torch.jit.load(pt_path).to(self.test_device)
+        self.last_w = 416
+        self.last_h = 416
+        self.grids = None
+
+    @torch.no_grad()
+    def detect(self, srcimg):
+        # t0=time.time()
+
+        h0, w0 = srcimg.shape[:2]  # orig hw
+        r = self.inpSize / min(h0, w0)  # resize image to img_size
+        h1 = int(h0*r+31)//32*32
+        w1 = int(w0*r+31)//32*32
+
+        img = cv2.resize(srcimg, (w1,h1), interpolation=cv2.INTER_LINEAR)
+
+        # Convert
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR to RGB
+
+        # Run inference
+        img = torch.from_numpy(img).to(self.test_device).permute(2,0,1)
+        img = img.float()/255  # uint8 to fp16/32  0-1
+        if img.ndimension() == 3:
+            img = img.unsqueeze(0)
+
+        # Inference
+        if h1 != self.last_h or w1 != self.last_w or self.grids is None:
+            grids = []
+            for scale in [8,16,32]:
+                ny = h1//scale
+                nx = w1//scale
+                yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
+                grid = torch.stack((xv, yv), 2).view((1,1,ny, nx, 2)).float()
+                grids.append(grid.to(self.test_device))
+            self.grids = grids
+            self.last_w = w1
+            self.last_h = h1
+
+        pred = self.model(img, self.grids).cpu()
+
+        # Apply NMS
+        det = non_max_suppression_face(pred, self.conf_thres, self.iou_thres)[0]
+        # Process detections
+        # det = pred[0]
+        bboxes = np.zeros((det.shape[0], 4))
+        kpss = np.zeros((det.shape[0], 5, 2))
+        scores = np.zeros((det.shape[0]))
+        # gn = torch.tensor([w0, h0, w0, h0]).to(pred)  # normalization gain whwh
+        # gn_lks = torch.tensor([w0, h0, w0, h0, w0, h0, w0, h0, w0, h0]).to(pred)  # normalization gain landmarks
+        det = det.cpu().numpy()
+
+        for j in range(det.shape[0]):
+            # xywh = (xyxy2xywh(det[j, :4].view(1, 4)) / gn).view(4).cpu().numpy()
+            bboxes[j, 0] = det[j, 0] * w0/w1
+            bboxes[j, 1] = det[j, 1] * h0/h1
+            bboxes[j, 2] = det[j, 2] * w0/w1 - bboxes[j, 0]
+            bboxes[j, 3] = det[j, 3] * h0/h1 - bboxes[j, 1]
+            scores[j] = det[j, 4]
+            # landmarks = (det[j, 5:15].view(1, 10) / gn_lks).view(5,2).cpu().numpy()
+            kpss[j, :, :] = det[j, 5:15].reshape(5, 2) * np.array([[w0/w1,h0/h1]])
+                # class_num = det[j, 15].cpu().numpy()
+                # orgimg = show_results(orgimg, xywh, conf, landmarks, class_num)
+        return bboxes, kpss, scores

hymm_sp/data_kits/ffmpeg_utils.py

@@ -0,0 +1,184 @@
+import skvideo
+# assert skvideo.__version__ >= "1.1.11"
+import os
+
+import skvideo.io
+import cv2
+
+# install the following packages: #
+# conda install -c conda-forge scikit-video ffmpeg  #
+import os
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+from einops import rearrange
+
+
+
+class VideoUtils(object):
+    def __init__(self, video_path=None, output_video_path=None, bit_rate='origin', fps=25):
+        if video_path is not None:
+            meta_data = skvideo.io.ffprobe(video_path)
+            # avg_frame_rate = meta_data['video']['@r_frame_rate']
+            # a, b = avg_frame_rate.split('/')
+            # fps = float(a) / float(b)
+            # fps = 25
+            codec_name = 'libx264'
+            # codec_name = meta_data['video'].get('@codec_name')
+            # if codec_name=='hevc':
+            #     codec_name='h264'
+            # profile = meta_data['video'].get('@profile')
+            color_space = meta_data['video'].get('@color_space')
+            color_transfer = meta_data['video'].get('@color_transfer')
+            color_primaries = meta_data['video'].get('@color_primaries')
+            color_range = meta_data['video'].get('@color_range')
+            pix_fmt = meta_data['video'].get('@pix_fmt')
+            if bit_rate=='origin':
+                bit_rate = meta_data['video'].get('@bit_rate')
+            else:
+                bit_rate=None
+            if pix_fmt is None:
+                pix_fmt = 'yuv420p'
+
+            reader_output_dict = {'-r': str(fps)}
+            writer_input_dict = {'-r': str(fps)}
+            writer_output_dict = {'-pix_fmt': pix_fmt, '-r': str(fps), '-vcodec':str(codec_name)}
+            # if bit_rate is not None:
+            #     writer_output_dict['-b:v'] = bit_rate
+            writer_output_dict['-crf'] = '17'
+
+            # if video has alpha channel, convert to bgra, uint16 to process
+            if pix_fmt.startswith('yuva'):
+                writer_input_dict['-pix_fmt'] = 'bgra64le'
+                reader_output_dict['-pix_fmt'] = 'bgra64le'
+            elif pix_fmt.endswith('le'):
+                writer_input_dict['-pix_fmt'] = 'bgr48le'
+                reader_output_dict['-pix_fmt'] = 'bgr48le'
+            else:
+                writer_input_dict['-pix_fmt'] = 'bgr24'
+                reader_output_dict['-pix_fmt'] = 'bgr24'
+
+            if color_range is not None:
+                writer_output_dict['-color_range'] = color_range
+                writer_input_dict['-color_range'] = color_range
+            if color_space is not None:
+                writer_output_dict['-colorspace'] = color_space
+                writer_input_dict['-colorspace'] = color_space
+            if color_primaries is not None:
+                writer_output_dict['-color_primaries'] = color_primaries
+                writer_input_dict['-color_primaries'] = color_primaries
+            if color_transfer is not None:
+                writer_output_dict['-color_trc'] = color_transfer
+                writer_input_dict['-color_trc'] = color_transfer
+
+            writer_output_dict['-sws_flags'] = 'full_chroma_int+bitexact+accurate_rnd'
+            reader_output_dict['-sws_flags'] = 'full_chroma_int+bitexact+accurate_rnd'
+            # writer_input_dict['-pix_fmt'] = 'bgr48le'
+            # reader_output_dict = {'-pix_fmt': 'bgr48le'}
+
+            # -s 1920x1080
+            # writer_input_dict['-s'] = '1920x1080'
+            # writer_output_dict['-s'] = '1920x1080'
+            # writer_input_dict['-s'] = '1080x1920'
+            # writer_output_dict['-s'] = '1080x1920'
+
+            print(writer_input_dict)
+            print(writer_output_dict)
+
+            self.reader = skvideo.io.FFmpegReader(video_path, outputdict=reader_output_dict)
+        else:
+            
+            # fps = 25
+            codec_name = 'libx264'
+            bit_rate=None
+            pix_fmt = 'yuv420p'
+
+            reader_output_dict = {'-r': str(fps)}
+            writer_input_dict = {'-r': str(fps)}
+            writer_output_dict = {'-pix_fmt': pix_fmt, '-r': str(fps), '-vcodec':str(codec_name)}
+            # if bit_rate is not None:
+            #     writer_output_dict['-b:v'] = bit_rate
+            writer_output_dict['-crf'] = '17'
+
+            # if video has alpha channel, convert to bgra, uint16 to process
+            if pix_fmt.startswith('yuva'):
+                writer_input_dict['-pix_fmt'] = 'bgra64le'
+                reader_output_dict['-pix_fmt'] = 'bgra64le'
+            elif pix_fmt.endswith('le'):
+                writer_input_dict['-pix_fmt'] = 'bgr48le'
+                reader_output_dict['-pix_fmt'] = 'bgr48le'
+            else:
+                writer_input_dict['-pix_fmt'] = 'bgr24'
+                reader_output_dict['-pix_fmt'] = 'bgr24'
+
+            writer_output_dict['-sws_flags'] = 'full_chroma_int+bitexact+accurate_rnd'
+            print(writer_input_dict)
+            print(writer_output_dict)
+
+        if output_video_path is not None:
+            self.writer = skvideo.io.FFmpegWriter(output_video_path, inputdict=writer_input_dict, outputdict=writer_output_dict, verbosity=1)
+
+    def getframes(self):
+        return self.reader.nextFrame()
+
+    def writeframe(self, frame):
+        if frame is None:
+            self.writer.close()
+        else:
+            self.writer.writeFrame(frame)
+
+
+def save_videos_from_pil(pil_images, path, fps=8):
+    save_fmt = ".mp4"
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    width, height = pil_images[0].size
+
+    if save_fmt == ".mp4":
+        video_cap = VideoUtils(output_video_path=path, fps=fps)
+        for pil_image in pil_images:
+            image_cv2 = np.array(pil_image)[:,:,[2,1,0]]
+            video_cap.writeframe(image_cv2)
+        video_cap.writeframe(None)
+
+    elif save_fmt == ".gif":
+        pil_images[0].save(
+            fp=path,
+            format="GIF",
+            append_images=pil_images[1:],
+            save_all=True,
+            duration=(1 / fps * 1000),
+            loop=0,
+            optimize=False,
+            lossless=True
+        )
+    else:
+        raise ValueError("Unsupported file type. Use .mp4 or .gif.")
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    height, width = videos.shape[-2:]
+    outputs = []
+
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)  # (c h w)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)  # (h w c)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = (x * 255).numpy().astype(np.uint8)
+        x = Image.fromarray(x)
+
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    save_videos_from_pil(outputs, path, fps)
+    
+def save_video(video, path: str, rescale=False, n_rows=6, fps=8):
+    outputs = []
+    for x in video:
+        x = Image.fromarray(x)
+        outputs.append(x)
+    
+    save_videos_from_pil(outputs, path, fps)

hymm_sp/diffusion/__init__.py

@@ -0,0 +1,30 @@
+from .pipelines import HunyuanVideoAudioPipeline
+from .schedulers import FlowMatchDiscreteScheduler
+
+
+def load_diffusion_pipeline(args, rank, vae, text_encoder, text_encoder_2, model, scheduler=None,
+                            device=None, progress_bar_config=None):
+    """ Load the denoising scheduler for inference. """
+    if scheduler is None:
+        scheduler = FlowMatchDiscreteScheduler(shift=args.flow_shift_eval_video, reverse=args.flow_reverse, solver=args.flow_solver, )
+        
+    # Only enable progress bar for rank 0
+    progress_bar_config = progress_bar_config or {'leave': True, 'disable': rank != 0}
+
+    pipeline = HunyuanVideoAudioPipeline(vae=vae,
+                                       text_encoder=text_encoder,
+                                       text_encoder_2=text_encoder_2,
+                                       transformer=model,
+                                       scheduler=scheduler,
+                                    #    safety_checker=None,
+                                    #    feature_extractor=None,
+                                    #    requires_safety_checker=False,
+                                       progress_bar_config=progress_bar_config,
+                                       args=args,
+                                       )
+    if args.cpu_offload:   #   avoid oom
+        pass
+    else:
+        pipeline = pipeline.to(device)
+
+    return pipeline

hymm_sp/diffusion/pipelines/__init__.py

@@ -0,0 +1 @@
+from .pipeline_hunyuan_video_audio import HunyuanVideoAudioPipeline

hymm_sp/diffusion/pipelines/pipeline_hunyuan_video_audio.py

@@ -0,0 +1,1363 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+import numpy as np
+import torch
+from packaging import version
+from diffusers.utils import BaseOutput
+from dataclasses import dataclass
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, ImageProjection
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    deprecate,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+from hymm_sp.constants import PRECISION_TO_TYPE
+from hymm_sp.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+from hymm_sp.text_encoder import TextEncoder
+from einops import rearrange
+from ...modules import HYVideoDiffusionTransformer
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """"""
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+@dataclass
+class HunyuanVideoPipelineOutput(BaseOutput):
+    videos: Union[torch.Tensor, np.ndarray]
+
+
+class HunyuanVideoAudioPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-video generation using HunyuanVideo.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        text_encoder ([`TextEncoder`]):
+            Frozen text-encoder.
+        text_encoder_2 ([`TextEncoder`]):
+            Frozen text-encoder_2.
+        transformer ([`HYVideoDiffusionTransformer`]):
+            A `HYVideoDiffusionTransformer` to denoise the encoded video latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = ["text_encoder_2"]
+    _exclude_from_cpu_offload = ["transformer"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: TextEncoder,
+        transformer: HYVideoDiffusionTransformer,
+        scheduler: KarrasDiffusionSchedulers,
+        text_encoder_2: Optional[TextEncoder] = None,
+        progress_bar_config: Dict[str, Any] = None,
+        args=None,
+    ):
+        super().__init__()
+
+        # ==========================================================================================
+        if progress_bar_config is None:
+            progress_bar_config = {}
+        if not hasattr(self, '_progress_bar_config'):
+            self._progress_bar_config = {}
+        self._progress_bar_config.update(progress_bar_config)
+
+        self.args = args
+        # ==========================================================================================
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            transformer=transformer,
+            scheduler=scheduler,
+            text_encoder_2=text_encoder_2
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        
+    def encode_prompt(
+        self,
+        prompt,
+        name,
+        device,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        pixel_value_llava: Optional[torch.Tensor] = None,
+        uncond_pixel_value_llava: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        lora_scale: Optional[float] = None,
+        clip_skip: Optional[int] = None,
+        text_encoder: Optional[TextEncoder] = None,
+        data_type: Optional[str] = "image",
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_videos_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            pixel_value_llava (`torch.Tensor`, *optional*):
+                The image tensor for llava. 
+            uncond_pixel_value_llava (`torch.Tensor`, *optional*):
+                The image tensor for llava.  Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            attention_mask (`torch.Tensor`, *optional*):
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_attention_mask (`torch.Tensor`, *optional*):
+            lora_scale (`float`, *optional*):
+                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            text_encoder (TextEncoder, *optional*):
+        """
+        if text_encoder is None:
+            text_encoder = self.text_encoder
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if not USE_PEFT_BACKEND:
+                adjust_lora_scale_text_encoder(text_encoder.model, lora_scale)
+            else:
+                scale_lora_layers(text_encoder.model, lora_scale)
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: process multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer)
+            text_inputs = text_encoder.text2tokens(prompt, data_type=data_type, name=name)
+
+            if pixel_value_llava is not None:
+                text_inputs['pixel_value_llava'] = pixel_value_llava
+                text_inputs['attention_mask'] = torch.cat([text_inputs['attention_mask'], torch.ones((1, 575 * len(pixel_value_llava))).to(text_inputs['attention_mask'])], dim=1)
+
+            if clip_skip is None:
+                prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type)
+                prompt_embeds = prompt_outputs.hidden_state
+            else:
+                prompt_outputs = text_encoder.encode(text_inputs, output_hidden_states=True, data_type=data_type)
+                # Access the `hidden_states` first, that contains a tuple of
+                # all the hidden states from the encoder layers. Then index into
+                # the tuple to access the hidden states from the desired layer.
+                prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)]
+                # We also need to apply the final LayerNorm here to not mess with the
+                # representations. The `last_hidden_states` that we typically use for
+                # obtaining the final prompt representations passes through the LayerNorm
+                # layer.
+                prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds)
+
+            attention_mask = prompt_outputs.attention_mask
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+                bs_embed, seq_len = attention_mask.shape
+                attention_mask = attention_mask.repeat(1, num_videos_per_prompt)
+                attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len)
+
+        if text_encoder is not None:
+            prompt_embeds_dtype = text_encoder.dtype
+        elif self.transformer is not None:
+            prompt_embeds_dtype = self.transformer.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        if prompt_embeds.ndim == 2:
+            bs_embed, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1)
+        else:
+            bs_embed, seq_len, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            # textual inversion: process multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, text_encoder.tokenizer)            
+            uncond_input = text_encoder.text2tokens(uncond_tokens, data_type=data_type)
+            if uncond_pixel_value_llava is not None:
+                uncond_input['pixel_value_llava'] = uncond_pixel_value_llava
+                uncond_input['attention_mask'] = torch.cat([uncond_input['attention_mask'], torch.ones((1, 575 * len(uncond_pixel_value_llava))).to(uncond_input['attention_mask'])], dim=1)
+
+            negative_prompt_outputs = text_encoder.encode(uncond_input, data_type=data_type)
+            negative_prompt_embeds = negative_prompt_outputs.hidden_state
+
+            negative_attention_mask = negative_prompt_outputs.attention_mask
+            if negative_attention_mask is not None:
+                negative_attention_mask = negative_attention_mask.to(device)
+                _, seq_len = negative_attention_mask.shape
+                negative_attention_mask = negative_attention_mask.repeat(1, num_videos_per_prompt)
+                negative_attention_mask = negative_attention_mask.view(batch_size * num_videos_per_prompt, seq_len)
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            if negative_prompt_embeds.ndim == 2:
+                negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt)
+                negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, -1)
+            else:
+                negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+                negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        if text_encoder is not None:
+            if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(text_encoder.model, lora_scale)
+
+        return prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask
+
+    def encode_prompt_audio_text_base(
+        self, 
+        prompt, 
+        uncond_prompt, 
+        pixel_value_llava, 
+        uncond_pixel_value_llava, 
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        lora_scale: Optional[float] = None,
+        clip_skip: Optional[int] = None,
+        text_encoder: Optional[TextEncoder] = None,
+        data_type: Optional[str] = "image",
+    ):
+        if text_encoder is None:
+            text_encoder = self.text_encoder
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if not USE_PEFT_BACKEND:
+                adjust_lora_scale_text_encoder(text_encoder.model, lora_scale)
+            else:
+                scale_lora_layers(text_encoder.model, lora_scale)
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+            
+        prompt_embeds = None
+        
+        if prompt_embeds is None:
+            # textual inversion: process multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer)
+            text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) # data_type: video, text_inputs: {'input_ids', 'attention_mask'}
+            
+            text_keys = ['input_ids', 'attention_mask']
+            
+            if pixel_value_llava is not None:
+                text_inputs['pixel_value_llava'] = pixel_value_llava
+                text_inputs['attention_mask'] = torch.cat([text_inputs['attention_mask'], torch.ones((1, 575)).to(text_inputs['attention_mask'])], dim=1)
+
+        
+            if clip_skip is None:
+                prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type)
+                prompt_embeds = prompt_outputs.hidden_state
+            else:
+                prompt_outputs = text_encoder.encode(text_inputs, output_hidden_states=True, data_type=data_type)
+                # Access the `hidden_states` first, that contains a tuple of
+                # all the hidden states from the encoder layers. Then index into
+                # the tuple to access the hidden states from the desired layer.
+                prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)]
+                # We also need to apply the final LayerNorm here to not mess with the
+                # representations. The `last_hidden_states` that we typically use for
+                # obtaining the final prompt representations passes through the LayerNorm
+                # layer.
+                prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds)
+
+            attention_mask = prompt_outputs.attention_mask
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+                bs_embed, seq_len = attention_mask.shape
+                attention_mask = attention_mask.repeat(1, num_images_per_prompt)
+                attention_mask = attention_mask.view(bs_embed * num_images_per_prompt, seq_len)
+
+        if text_encoder is not None:
+            prompt_embeds_dtype = text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        if prompt_embeds.ndim == 2:
+            bs_embed, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, -1)
+        else:
+            bs_embed, seq_len, _ = prompt_embeds.shape
+            # duplicate text embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            # textual inversion: process multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, text_encoder.tokenizer)            
+            # max_length = prompt_embeds.shape[1]
+            uncond_input = text_encoder.text2tokens(uncond_tokens, data_type=data_type)
+
+            # if hasattr(text_encoder.model.config, "use_attention_mask") and text_encoder.model.config.use_attention_mask:
+            #     attention_mask = uncond_input.attention_mask.to(device)
+            # else:
+            #     attention_mask = None
+            if uncond_pixel_value_llava is not None:
+                uncond_input['pixel_value_llava'] = uncond_pixel_value_llava
+                uncond_input['attention_mask'] = torch.cat([uncond_input['attention_mask'], torch.ones((1, 575)).to(uncond_input['attention_mask'])], dim=1)
+
+            negative_prompt_outputs = text_encoder.encode(uncond_input, data_type=data_type)
+            negative_prompt_embeds = negative_prompt_outputs.hidden_state
+
+            negative_attention_mask = negative_prompt_outputs.attention_mask
+            if negative_attention_mask is not None:
+                negative_attention_mask = negative_attention_mask.to(device)
+                _, seq_len = negative_attention_mask.shape
+                negative_attention_mask = negative_attention_mask.repeat(1, num_images_per_prompt)
+                negative_attention_mask = negative_attention_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            if negative_prompt_embeds.ndim == 2:
+                negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt)
+                negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+            else:
+                negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+                negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        if text_encoder is not None:
+            if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(text_encoder.model, lora_scale)
+
+        return prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask
+
+    def decode_latents(self, latents, enable_tiling=True):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        if enable_tiling:
+            self.vae.enable_tiling()
+            image = self.vae.decode(latents, return_dict=False)[0]
+            self.vae.disable_tiling()
+        else:
+            image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        if image.ndim==4: image = image.cpu().permute(0, 2, 3, 1).float()
+        else: image = image.cpu().float()
+        return image
+
+    def prepare_extra_func_kwargs(self, func, kwargs):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+        extra_step_kwargs = {}
+
+        for k, v in kwargs.items():
+            accepts = k in set(inspect.signature(func).parameters.keys())
+            if accepts:
+                extra_step_kwargs[k] = v
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        frame,
+        callback_steps,
+        pixel_value_llava=None,
+        uncond_pixel_value_llava=None,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        vae_ver='88-4c-sd'
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if frame is not None:
+            if '884' in vae_ver:
+                if frame!=1 and (frame-1)%4!=0:
+                    raise ValueError(f'`frame` has to be 1 or a multiple of 4 but is {frame}.')
+            elif '888' in vae_ver:
+                if frame!=1 and (frame-1)%8!=0:
+                    raise ValueError(f'`frame` has to be 1 or a multiple of 8 but is {frame}.')
+
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if pixel_value_llava is not None and uncond_pixel_value_llava is not None:
+            if len(pixel_value_llava) != len(uncond_pixel_value_llava):
+                raise ValueError(
+                    "`pixel_value_llava` and `uncond_pixel_value_llava` must have the same length when passed directly, but"
+                    f" got: `pixel_value_llava` {len(pixel_value_llava)} != `uncond_pixel_value_llava`"
+                    f" {len(uncond_pixel_value_llava)}."
+                )
+            
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def get_timesteps(self, num_inference_steps, strength, device):
+        # get the original timestep using init_timestep
+        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+        t_start = max(num_inference_steps - init_timestep, 0)
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(t_start * self.scheduler.order)
+
+        return timesteps.to(device), num_inference_steps - t_start
+    
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, frame, dtype, device, generator, latents=None, ref_latents=None, timestep=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            frame,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+
+        if timestep is not None:
+            init_latents = ref_latents.clone().repeat(1,1,frame,1,1).to(device).to(dtype)
+            latents = latents
+
+        # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler
+        if hasattr(self.scheduler, "init_noise_sigma"):
+            latents = latents * self.scheduler.init_noise_sigma
+
+        return latents
+
+    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+    def get_guidance_scale_embedding(
+        self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
+    ) -> torch.Tensor:
+        """
+        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+        Args:
+            w (`torch.Tensor`):
+                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+            embedding_dim (`int`, *optional*, defaults to 512):
+                Dimension of the embeddings to generate.
+            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+                Data type of the generated embeddings.
+
+        Returns:
+            `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+        """
+        assert len(w.shape) == 1
+        w = w * 1000.0
+
+        half_dim = embedding_dim // 2
+        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+        emb = w.to(dtype)[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+        if embedding_dim % 2 == 1:  # zero pad
+            emb = torch.nn.functional.pad(emb, (0, 1))
+        assert emb.shape == (w.shape[0], embedding_dim)
+        return emb
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def guidance_rescale(self):
+        return self._guidance_rescale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None
+        return self._guidance_scale > 1
+
+    @property
+    def cross_attention_kwargs(self):
+        return self._cross_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],                              
+ 
+        ref_latents: Union[torch.Tensor],                            # [1, 16, 1, h//8, w//8]
+        uncond_ref_latents: Union[torch.Tensor],
+        pixel_value_llava: Union[torch.Tensor],                # [1, 3, 336, 336]
+        uncond_pixel_value_llava: Union[torch.Tensor],
+        face_masks: Union[torch.Tensor],                              # [b f h w]
+        audio_prompts: Union[torch.Tensor], 
+        uncond_audio_prompts: Union[torch.Tensor], 
+        motion_exp: Union[torch.Tensor], 
+        motion_pose: Union[torch.Tensor], 
+        fps: Union[torch.Tensor],
+      
+        height: int,
+        width: int,
+        frame: int,
+        data_type: str = "video",
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_attention_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[
+            Union[
+                Callable[[int, int, Dict], None],
+                PipelineCallback,
+                MultiPipelineCallbacks,
+            ]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
+        vae_ver: str = "88-4c-sd",
+        enable_tiling: bool = False,
+        n_tokens: Optional[int] = None,
+        embedded_guidance_scale: Optional[float] = None,
+        **kwargs,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            height (`int`):
+                The height in pixels of the generated image.
+            width (`int`):
+                The width in pixels of the generated image.
+            video_length (`int`):
+                The number of frames in the generated video.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+                
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a
+                plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.0):
+                Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
+                using zero terminal SNR.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~HunyuanVideoPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list with the generated images and the
+                second element is a list of `bool`s indicating whether the corresponding generated image contains
+                "not-safe-for-work" (nsfw) content.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        cpu_offload = kwargs.get("cpu_offload", 0)
+
+        # 0. Default height and width to transformer
+        # height = height or self.transformer.config.sample_size * self.vae_scale_factor
+        # width = width or self.transformer.config.sample_size * self.vae_scale_factor
+        # to deal with lora scaling and other possible forward hooks
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            frame,
+            callback_steps,
+            pixel_value_llava,
+            uncond_pixel_value_llava,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            callback_on_step_end_tensor_inputs,
+            vae_ver=vae_ver
+        )
+
+        self._guidance_scale = guidance_scale
+        self.start_cfg_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+        )
+
+ 
+        # ========== Encode text prompt (image prompt) ==========
+        prompt_embeds, negative_prompt_embeds, prompt_mask, negative_prompt_mask = \
+            self.encode_prompt_audio_text_base(
+                prompt=prompt,
+                uncond_prompt=negative_prompt,
+                pixel_value_llava=pixel_value_llava,
+                uncond_pixel_value_llava=uncond_pixel_value_llava,
+                device=device,
+                num_images_per_prompt=num_videos_per_prompt,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+                negative_prompt=negative_prompt,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                lora_scale=lora_scale,
+                clip_skip=self.clip_skip,
+                text_encoder=self.text_encoder,
+                data_type=data_type, 
+                # **kwargs
+            )
+        if self.text_encoder_2 is not None:
+            prompt_embeds_2, negative_prompt_embeds_2, prompt_mask_2, negative_prompt_mask_2 = \
+                self.encode_prompt_audio_text_base(
+                    prompt=prompt,
+                    uncond_prompt=negative_prompt,
+                    pixel_value_llava=None,
+                    uncond_pixel_value_llava=None,
+                    device=device,
+                    num_images_per_prompt=num_videos_per_prompt,
+                    do_classifier_free_guidance=self.do_classifier_free_guidance,
+                    negative_prompt=negative_prompt,
+                    prompt_embeds=None,
+                    negative_prompt_embeds=None,
+                    lora_scale=lora_scale,
+                    clip_skip=self.clip_skip,
+                    text_encoder=self.text_encoder_2,
+                    # **kwargs
+                )
+        else:
+            prompt_embeds_2 = None
+            negative_prompt_embeds_2 = None
+            prompt_mask_2 = None
+            negative_prompt_mask_2 = None
+
+
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if self.do_classifier_free_guidance:
+            prompt_embeds_input = torch.cat([negative_prompt_embeds, prompt_embeds])
+            if prompt_mask is not None:
+                prompt_mask_input = torch.cat([negative_prompt_mask, prompt_mask])
+            if prompt_embeds_2 is not None:
+                prompt_embeds_2_input = torch.cat([negative_prompt_embeds_2, prompt_embeds_2])
+            if prompt_mask_2 is not None:
+                prompt_mask_2_input = torch.cat([negative_prompt_mask_2, prompt_mask_2])
+            
+        if self.do_classifier_free_guidance:
+            ref_latents = torch.cat([ref_latents, ref_latents], dim=0)
+
+
+        # 4. Prepare timesteps
+        extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.set_timesteps, {"n_tokens": n_tokens}
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas, **extra_set_timesteps_kwargs,
+        )
+
+        video_length = audio_prompts.shape[1] // 4 * 4 + 1
+        if "884" in vae_ver:
+            video_length = (video_length - 1) // 4 + 1
+        elif "888" in vae_ver:
+            video_length = (video_length - 1) // 8 + 1
+        else:
+            video_length = video_length
+
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        infer_length = (audio_prompts.shape[1] // 128 + 1) * 32 + 1
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            infer_length,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+            ref_latents[-1:],
+            timesteps[:1]
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step, {"generator": generator, "eta": eta},
+        )
+
+        target_dtype = PRECISION_TO_TYPE[self.args.precision]
+        autocast_enabled = (target_dtype != torch.float32) and not self.args.val_disable_autocast
+        vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision]
+        vae_autocast_enabled = (vae_dtype != torch.float32) and not self.args.val_disable_autocast
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        latents_all = latents.clone()
+        pad_audio_length = (audio_prompts.shape[1] // 128 + 1) * 128 + 4 - audio_prompts.shape[1]
+        audio_prompts_all = torch.cat([audio_prompts, torch.zeros_like(audio_prompts[:, :pad_audio_length])], dim=1)
+
+
+        shift = 0
+        shift_offset = 10
+        frames_per_batch = 33
+        self.cache_tensor = None
+
+        """ If the total length is shorter than 129, shift is not required """
+        if video_length == 33 or infer_length == 33:
+            infer_length = 33
+            shift_offset = 0
+            latents_all = latents_all[:, :, :33]
+            audio_prompts_all = audio_prompts_all[:, :132]
+
+        if cpu_offload: torch.cuda.empty_cache()
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # init 
+                pred_latents = torch.zeros_like(
+                    latents_all,
+                    dtype=latents_all.dtype,
+                )
+                counter = torch.zeros(
+                    (latents_all.shape[0], latents_all.shape[1], infer_length, 1, 1),
+                    dtype=latents_all.dtype,
+                ).to(device=latents_all.device)
+
+                for index_start in range(0, infer_length, frames_per_batch):
+                    self.scheduler._step_index = None
+
+                    index_start = index_start - shift
+                    
+                    idx_list = [ii % latents_all.shape[2] for ii in range(index_start, index_start + frames_per_batch)]
+                    latents = latents_all[:, :, idx_list].clone()
+
+                    idx_list_audio = [ii % audio_prompts_all.shape[1] for ii in range(index_start * 4, (index_start + frames_per_batch) * 4 - 3)]
+                    audio_prompts = audio_prompts_all[:, idx_list_audio].clone()
+
+                    # expand the latents if we are doing classifier free guidance
+                    if self.do_classifier_free_guidance:
+                        latent_model_input = torch.cat([latents] * 2)
+                    else:
+                        latent_model_input = latents 
+                        
+                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                    
+                    if self.do_classifier_free_guidance:
+                        if i < 10:
+                            self._guidance_scale = (1 - i / len(timesteps)) * (self.start_cfg_scale - 2) + 2
+                            audio_prompts_input = torch.cat([uncond_audio_prompts, audio_prompts], dim=0)
+                            face_masks_input = torch.cat([face_masks * 0.6] * 2, dim=0)
+                        else:
+                            # define 10-50 step cfg
+                            self._guidance_scale = (1 - i / len(timesteps)) * (6.5 - 3.5) + 3.5  # 5-2 +2
+
+                            prompt_embeds_input = torch.cat([prompt_embeds, prompt_embeds])
+                            if prompt_mask is not None:
+                                prompt_mask_input = torch.cat([prompt_mask, prompt_mask])
+                            if prompt_embeds_2 is not None:
+                                prompt_embeds_2_input = torch.cat([prompt_embeds_2, prompt_embeds_2])
+                            if prompt_mask_2 is not None:
+                                prompt_mask_2_input = torch.cat([prompt_mask_2, prompt_mask_2])
+                            audio_prompts_input = torch.cat([uncond_audio_prompts, audio_prompts], dim=0)
+                            face_masks_input = torch.cat([face_masks] * 2, dim=0)
+
+                        motion_exp_input = torch.cat([motion_exp] * 2, dim=0)
+                        motion_pose_input = torch.cat([motion_pose] * 2, dim=0)
+                        fps_input = torch.cat([fps] * 2, dim=0)
+    
+                    else:
+                        audio_prompts_input = audio_prompts
+                        face_masks_input = face_masks
+                        motion_exp_input = motion_exp
+                        motion_pose_input = motion_pose
+                        fps_input = fps
+
+                    t_expand = t.repeat(latent_model_input.shape[0])
+                    guidance_expand = None
+
+                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                        
+                        no_cache_steps = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] + list(range(15, 42, 5)) + [41, 42, 43, 44, 45, 46, 47, 48, 49]
+                        img_len = (latent_model_input.shape[-1] // 2)  * (latent_model_input.shape[-2] // 2) * latent_model_input.shape[-3]
+                        img_ref_len = (latent_model_input.shape[-1] // 2)  * (latent_model_input.shape[-2] // 2) * (latent_model_input.shape[-3]+1) 
+                        if i in no_cache_steps:
+                            is_cache = False
+
+                            if latent_model_input.shape[-1]*latent_model_input.shape[-2]>64*112 and cpu_offload:
+                                if i==0:
+                                    print(f'cpu_offload={cpu_offload} and {latent_model_input.shape[-2:]} is large, split infer noise-pred')
+                                
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input[:1],        
+                                    "motion_pose": motion_pose_input[:1],      
+                                    "fps": fps_input[:1],                       
+                                    "audio_prompts": audio_prompts_input[:1],  
+                                    "face_mask": face_masks_input[:1]        
+                                }
+                                noise_pred_uncond = self.transformer(latent_model_input[:1], t_expand[:1], ref_latents=ref_latents[:1], text_states=prompt_embeds_input[:1], text_mask=prompt_mask_input[:1], text_states_2=prompt_embeds_2_input[:1], freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+                                uncond_cache_tensor = self.transformer.cache_out
+                                torch.cuda.empty_cache()
+
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input[1:],        
+                                    "motion_pose": motion_pose_input[1:],      
+                                    "fps": fps_input[1:],                       
+                                    "audio_prompts": audio_prompts_input[1:],  
+                                    "face_mask": face_masks_input[1:]        
+                                }
+                                noise_pred_text = self.transformer(latent_model_input[1:], t_expand[1:], ref_latents=ref_latents[1:], text_states=prompt_embeds_input[1:], text_mask=prompt_mask_input[1:], text_states_2=prompt_embeds_2_input[1:], freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+                                self.transformer.cache_out = torch.cat([uncond_cache_tensor, self.transformer.cache_out], dim=0)
+
+                                noise_pred = torch.cat([noise_pred_uncond, noise_pred_text], dim=0)
+                                torch.cuda.empty_cache()
+                            else:
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input,        
+                                    "motion_pose": motion_pose_input,     
+                                    "fps": fps_input,                       
+                                    "audio_prompts": audio_prompts_input,   
+                                    "face_mask": face_masks_input         
+                                }
+                                noise_pred = self.transformer(latent_model_input, t_expand, ref_latents=ref_latents, text_states=prompt_embeds_input, text_mask=prompt_mask_input, text_states_2=prompt_embeds_2_input, freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+                                torch.cuda.empty_cache()
+
+                            if self.cache_tensor is None:
+                                self.cache_tensor = {
+                                    "ref": torch.zeros([latent_model_input.shape[0], latents_all.shape[-3], (latent_model_input.shape[-1] // 2)  * (latent_model_input.shape[-2] // 2), 3072]).to(self.transformer.cache_out.dtype).to(latent_model_input.device).clone(),
+                                    "img": torch.zeros([latent_model_input.shape[0], latents_all.shape[-3], (latent_model_input.shape[-1] // 2)  * (latent_model_input.shape[-2] // 2), 3072]).to(self.transformer.cache_out.dtype).to(latent_model_input.device).clone(),
+                                    "txt": torch.zeros([latent_model_input.shape[0], latents_all.shape[-3], prompt_embeds_input.shape[1], 3072]).to(self.transformer.cache_out.dtype).to(latent_model_input.device).clone(),
+                                }
+
+                            self.cache_tensor["ref"][:, idx_list] = self.transformer.cache_out[:, :img_ref_len-img_len].reshape(latent_model_input.shape[0], 1, -1, 3072).repeat(1, len(idx_list), 1, 1)
+                            self.cache_tensor["img"][:, idx_list] = self.transformer.cache_out[:, img_ref_len-img_len:img_ref_len].reshape(latent_model_input.shape[0], len(idx_list), -1, 3072)
+                            self.cache_tensor["txt"][:, idx_list] = self.transformer.cache_out[:, img_ref_len:].unsqueeze(1).repeat(1, len(idx_list), 1, 1)
+                            
+                        else:
+                            is_cache = True
+                            # self.transformer.cache_out[:, :img_ref_len-img_len] = self.cache_tensor["ref"][:, idx_list].mean(1)
+                            self.transformer.cache_out[:, :img_ref_len-img_len] = self.cache_tensor["ref"][:, idx_list][:, 0].clone()
+                            self.transformer.cache_out[:, img_ref_len-img_len:img_ref_len] = self.cache_tensor["img"][:, idx_list].reshape(-1, img_len, 3072).clone()
+                            self.transformer.cache_out[:, img_ref_len:] = self.cache_tensor["txt"][:, idx_list][:, 0].clone()
+                            
+                            if latent_model_input.shape[-1]*latent_model_input.shape[-2]>64*112 and cpu_offload:
+                                if i==0:
+                                    print(f'cpu_offload={cpu_offload} and {latent_model_input.shape[-2:]} is large, split infer noise-pred')
+                                
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input[:1],        
+                                    "motion_pose": motion_pose_input[:1],      
+                                    "fps": fps_input[:1],                       
+                                    "audio_prompts": audio_prompts_input[:1],  
+                                    "face_mask": face_masks_input[:1]        
+                                }
+                                tmp = self.transformer.cache_out.clone()
+                                self.transformer.cache_out = tmp[:1]
+                                noise_pred_uncond = self.transformer(latent_model_input[:1], t_expand[:1], ref_latents=ref_latents[:1], text_states=prompt_embeds_input[:1], text_mask=prompt_mask_input[:1], text_states_2=prompt_embeds_2_input[:1], freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+
+
+                                torch.cuda.empty_cache()
+
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input[1:],        
+                                    "motion_pose": motion_pose_input[1:],      
+                                    "fps": fps_input[1:],                       
+                                    "audio_prompts": audio_prompts_input[1:],  
+                                    "face_mask": face_masks_input[1:]        
+                                }
+                                self.transformer.cache_out = tmp[1:]
+                                noise_pred_text = self.transformer(latent_model_input[1:], t_expand[1:], ref_latents=ref_latents[1:], text_states=prompt_embeds_input[1:], text_mask=prompt_mask_input[1:], text_states_2=prompt_embeds_2_input[1:], freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+                                noise_pred = torch.cat([noise_pred_uncond, noise_pred_text], dim=0)
+
+                                self.transformer.cache_out = tmp
+                                torch.cuda.empty_cache()
+                            else:
+                                additional_kwargs = {
+                                    "motion_exp": motion_exp_input,        
+                                    "motion_pose": motion_pose_input,     
+                                    "fps": fps_input,                       
+                                    "audio_prompts": audio_prompts_input,   
+                                    "face_mask": face_masks_input         
+                                }
+                                noise_pred = self.transformer(latent_model_input, t_expand, ref_latents=ref_latents, text_states=prompt_embeds_input, text_mask=prompt_mask_input, text_states_2=prompt_embeds_2_input, freqs_cos=freqs_cis[0], freqs_sin=freqs_cis[1], guidance=guidance_expand, return_dict=True, is_cache=is_cache, **additional_kwargs,)['x']
+                                torch.cuda.empty_cache()
+                    # perform guidance
+                    if self.do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                      
+                    if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                        # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                        noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                    if callback_on_step_end is not None:
+                        callback_kwargs = {}
+                        for k in callback_on_step_end_tensor_inputs:
+                            callback_kwargs[k] = locals()[k]
+                        callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                        latents = callback_outputs.pop("latents", latents)
+                        prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                        negative_prompt_embeds = callback_outputs.pop(
+                            "negative_prompt_embeds", negative_prompt_embeds
+                        )
+                    latents = latents.to(torch.bfloat16)
+                    for iii in range(frames_per_batch):
+                        p = (index_start + iii) % pred_latents.shape[2]
+                        pred_latents[:, :, p] += latents[:, :, iii]
+                        counter[:, :, p] += 1
+
+                shift += shift_offset
+                shift = shift % frames_per_batch  
+                pred_latents  = pred_latents / counter
+                latents_all = pred_latents     
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    if progress_bar is not None:
+                        progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        latents = latents_all.float()[:, :, :video_length] 
+        if cpu_offload: torch.cuda.empty_cache()
+
+        if not output_type == "latent":
+            expand_temporal_dim = False
+            if len(latents.shape) == 4:
+                if isinstance(self.vae, AutoencoderKLCausal3D):
+                    latents = latents.unsqueeze(2)
+                    expand_temporal_dim = True
+            elif len(latents.shape) == 5:
+                pass
+            else:
+                raise ValueError(
+                    f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.")
+
+            if hasattr(self.vae.config, 'shift_factor') and self.vae.config.shift_factor:
+                latents = latents / self.vae.config.scaling_factor + self.vae.config.shift_factor
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
+                if enable_tiling:
+                    self.vae.enable_tiling()
+                    if cpu_offload:
+                        self.vae.post_quant_conv.to('cuda')
+                        self.vae.decoder.to('cuda')
+                    image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
+                    self.vae.disable_tiling()
+                    if cpu_offload:
+                        self.vae.post_quant_conv.to('cpu')
+                        self.vae.decoder.to('cpu')
+                        torch.cuda.empty_cache()
+                else:
+                    image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
+            if image is None:
+                return (None, )
+
+            if expand_temporal_dim or image.shape[2] == 1:
+                image = image.squeeze(2)
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        image = image.cpu().float()
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+        
+        if cpu_offload: torch.cuda.empty_cache()
+        if not return_dict:
+            return image
+        
+        return HunyuanVideoPipelineOutput(videos=image)

hymm_sp/diffusion/schedulers/__init__.py

@@ -0,0 +1 @@
+from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler