diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..2b8113889521e615669f771ab9fc45e6ffb7a03c 100644 --- a/.gitattributes +++ b/.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 diff --git a/assets/audio/2.WAV b/assets/audio/2.WAV new file mode 100644 index 0000000000000000000000000000000000000000..016714bab6bd8208d2714ece3f0007d09034a277 --- /dev/null +++ b/assets/audio/2.WAV @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:388d768354bb20f3aaa8327ef3391737c8150d3351fdaa04aa98b57caddc5dfb +size 3862572 diff --git a/assets/audio/3.WAV b/assets/audio/3.WAV new file mode 100644 index 0000000000000000000000000000000000000000..4025c9ed082d32834dfb2bb5d5db0c1caad9fcf1 --- /dev/null +++ b/assets/audio/3.WAV @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:854f2df163f4fc7bd69f09e8ca31758dedf8a56b191083d5fd4a5b25259b5fe2 +size 1921068 diff --git a/assets/audio/4.WAV b/assets/audio/4.WAV new file mode 100644 index 0000000000000000000000000000000000000000..0d1fe5987a5168b32ea51fb86a6ccb591a92fc29 --- /dev/null +++ b/assets/audio/4.WAV @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81d992d96a96829f27aef2d99fa63614b692cc9ed2cbac61e23e9a1dcc402b3a +size 3809324 diff --git a/assets/image/1.png b/assets/image/1.png new file mode 100644 index 0000000000000000000000000000000000000000..3934b7e4e6c54fd6688b61dac35071fd4864607d --- /dev/null +++ b/assets/image/1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d9bb2f06d84af5ca9347d684554f9bd8edb118749457b37bff6026a1b429843d +size 7359183 diff --git a/assets/image/2.png b/assets/image/2.png new file mode 100644 index 0000000000000000000000000000000000000000..3c1b7a2ca199c3f4780d2dacad9891d70a8b8f22 --- /dev/null +++ b/assets/image/2.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:97fdb135022e2d82366b5891b69d27b14fb183941c52cc6903acb5b6e74e89eb +size 5928730 diff --git a/assets/image/3.png b/assets/image/3.png new file mode 100644 index 0000000000000000000000000000000000000000..c2f5382c020e468600ac72aaabf491364b921f24 --- /dev/null +++ b/assets/image/3.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6616f4dde7435a504f11af1b8919124ae48a592858fdc1ab0ccacaa265f2271d +size 7343992 diff --git a/assets/image/4.png b/assets/image/4.png new file mode 100644 index 0000000000000000000000000000000000000000..d312b2ac07140a285def8590598c3e7e78ccab1b --- /dev/null +++ b/assets/image/4.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:931921def1c1c9336d1330554c1fa39b6090431dd720ccaa279e684b6253a800 +size 5276839 diff --git a/assets/image/src1.png b/assets/image/src1.png new file mode 100644 index 0000000000000000000000000000000000000000..ff2320b33af8e261fa9d46734ee293303ad5c036 --- /dev/null +++ b/assets/image/src1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c1ff64938aaa0950e33132e43cdc1e3e3a1cedc1b6ecdb33a067cb1971c4189 +size 10159438 diff --git a/assets/image/src2.png b/assets/image/src2.png new file mode 100644 index 0000000000000000000000000000000000000000..dbb04996d099b2e617822845b8e3385a1454f7e4 --- /dev/null +++ b/assets/image/src2.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4b4bb9edc4d01cec463d85d97cd1c9319e34d6f5d10e8d11d6c33a89f418b66 +size 9553059 diff --git a/assets/image/src3.png b/assets/image/src3.png new file mode 100644 index 0000000000000000000000000000000000000000..0f76d382f005034e474d28e2baad3fc799700bda --- /dev/null +++ b/assets/image/src3.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b51bed13f3fe69f5b98b95faa1e570a3600999bcc2479ef6080ac80f53255855 +size 9991903 diff --git a/assets/image/src4.png b/assets/image/src4.png new file mode 100644 index 0000000000000000000000000000000000000000..7d45025247e5a193469494d6940da5149cfda694 --- /dev/null +++ b/assets/image/src4.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:25b8bbf44d24bd25cb37495816cf1aa9cfcc9e32935dad074823bfdc590d0420 +size 9601751 diff --git a/assets/material/demo.png b/assets/material/demo.png new file mode 100644 index 0000000000000000000000000000000000000000..11f6a4daaaac8733505d0e147c18d79bc374f94f --- /dev/null +++ b/assets/material/demo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:32a0b751c12f40babc96f208cdfa3932aca74044c7689770b849e461e0b6aa1a +size 7675020 diff --git a/assets/material/logo.png b/assets/material/logo.png new file mode 100644 index 0000000000000000000000000000000000000000..803e7e155b28c1a3dd4330d95a879c94aca5bca3 --- /dev/null +++ b/assets/material/logo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2318f6ce82b18f91586050672e45cc637db61a33885de672b1b0d26e31b5de94 +size 161891 diff --git a/assets/material/method.png b/assets/material/method.png new file mode 100644 index 0000000000000000000000000000000000000000..35718a49263e5fe80c8fbdc26ea327a6090464c0 --- /dev/null +++ b/assets/material/method.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40ebe93eec368255ffaf58ad88f7aa5cdb605fd538f9272a3257803f9e86c338 +size 4942692 diff --git a/assets/material/teaser.png b/assets/material/teaser.png new file mode 100644 index 0000000000000000000000000000000000000000..76a551c02eb23adf29e7fc6c6801173d351ad45b --- /dev/null +++ b/assets/material/teaser.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8bbd88f126ad3a80ed613d90bb93e7a6bda346245df4aecd75b884ae00607ac4 +size 5434996 diff --git a/assets/test.csv b/assets/test.csv new file mode 100644 index 0000000000000000000000000000000000000000..a58c694894a0bcf4a63144d62acc410994359462 --- /dev/null +++ b/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 diff --git a/hymm_gradio/flask_audio.py b/hymm_gradio/flask_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..24ecc31cf3c6643b3d05146485c5c2dd1c95e976 --- /dev/null +++ b/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() + diff --git a/hymm_gradio/gradio_audio.py b/hymm_gradio/gradio_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..3100c8d55bc864028bebb84401dc39b64f4fe21b --- /dev/null +++ b/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_ = ''' +
+

Tencent HunyuanVideo-Avatar Demo

+
+ +''' +# 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) diff --git a/hymm_gradio/tool_for_end2end.py b/hymm_gradio/tool_for_end2end.py new file mode 100644 index 0000000000000000000000000000000000000000..4b1ba3f1d1781e2bf39fbbe1e5a95673a3227972 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/hymm_sp/__init__.py b/hymm_sp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/hymm_sp/__pycache__/__init__.cpython-310.pyc b/hymm_sp/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e71c0dd232c20c1c7118f1516996e125c66a190e Binary files /dev/null and b/hymm_sp/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/__pycache__/config.cpython-310.pyc b/hymm_sp/__pycache__/config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04736342535bbde1bf8d5594bd78fde8c0b674fe Binary files /dev/null and b/hymm_sp/__pycache__/config.cpython-310.pyc differ diff --git a/hymm_sp/__pycache__/constants.cpython-310.pyc b/hymm_sp/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0642332471e26d651188def1a2dec53ab8d8b562 Binary files /dev/null and b/hymm_sp/__pycache__/constants.cpython-310.pyc differ diff --git a/hymm_sp/__pycache__/helpers.cpython-310.pyc b/hymm_sp/__pycache__/helpers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba22320d88716614bd989a58ad7c352c08fd2c5e Binary files /dev/null and b/hymm_sp/__pycache__/helpers.cpython-310.pyc differ diff --git a/hymm_sp/__pycache__/inference.cpython-310.pyc b/hymm_sp/__pycache__/inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bc26fe2fe2ab810393782c22afc1684db576d62 Binary files /dev/null and b/hymm_sp/__pycache__/inference.cpython-310.pyc differ diff --git a/hymm_sp/__pycache__/sample_inference_audio.cpython-310.pyc b/hymm_sp/__pycache__/sample_inference_audio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73fae8bddc96b61f82f542a437a7bf7c25fb3dde Binary files /dev/null and b/hymm_sp/__pycache__/sample_inference_audio.cpython-310.pyc differ diff --git a/hymm_sp/config.py b/hymm_sp/config.py new file mode 100644 index 0000000000000000000000000000000000000000..949570b4745437b9bb110a49dbe07c9cf8bdd436 --- /dev/null +++ b/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 diff --git a/hymm_sp/constants.py b/hymm_sp/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9d52df91d43a894ab8a52e47d6edd3b4e10d39a7 --- /dev/null +++ b/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 +} diff --git a/hymm_sp/data_kits/__pycache__/audio_dataset.cpython-310.pyc b/hymm_sp/data_kits/__pycache__/audio_dataset.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..167a32ecd13fb60c8c7d8cb8f326822e83329e81 Binary files /dev/null and b/hymm_sp/data_kits/__pycache__/audio_dataset.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/__pycache__/audio_preprocessor.cpython-310.pyc b/hymm_sp/data_kits/__pycache__/audio_preprocessor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfcd9dbb22c29885e196f0f32ad3c39992fbda4c Binary files /dev/null and b/hymm_sp/data_kits/__pycache__/audio_preprocessor.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/__pycache__/data_tools.cpython-310.pyc b/hymm_sp/data_kits/__pycache__/data_tools.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f961ca00f4f4c14cb0d32d87b778680febe59fe Binary files /dev/null and b/hymm_sp/data_kits/__pycache__/data_tools.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/__pycache__/ffmpeg_utils.cpython-310.pyc b/hymm_sp/data_kits/__pycache__/ffmpeg_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07afed918376098aabe530cefeae7ecdcfc92d50 Binary files /dev/null and b/hymm_sp/data_kits/__pycache__/ffmpeg_utils.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/audio_dataset.py b/hymm_sp/data_kits/audio_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..e65cee54f0f836454ee70b524afd33a718080a47 --- /dev/null +++ b/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) + + + + \ No newline at end of file diff --git a/hymm_sp/data_kits/audio_preprocessor.py b/hymm_sp/data_kits/audio_preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..89ac4707c3d0b18f781746286fb540ab89ba6a62 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/hymm_sp/data_kits/data_tools.py b/hymm_sp/data_kits/data_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..a7d6077dbc79cba3e19281b3e4de7a1480158277 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/hymm_sp/data_kits/face_align/__init__.py b/hymm_sp/data_kits/face_align/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d6600221962b39f5270b363e29c7d257fd9830d9 --- /dev/null +++ b/hymm_sp/data_kits/face_align/__init__.py @@ -0,0 +1 @@ +from .align import AlignImage \ No newline at end of file diff --git a/hymm_sp/data_kits/face_align/__pycache__/__init__.cpython-310.pyc b/hymm_sp/data_kits/face_align/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..834523a54ad5249e4b862b7c9e2ad8caba3b475a Binary files /dev/null and b/hymm_sp/data_kits/face_align/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/face_align/__pycache__/align.cpython-310.pyc b/hymm_sp/data_kits/face_align/__pycache__/align.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74ac6f7859169c52439b0ac3bde1b40e6fedd9fc Binary files /dev/null and b/hymm_sp/data_kits/face_align/__pycache__/align.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/face_align/__pycache__/detface.cpython-310.pyc b/hymm_sp/data_kits/face_align/__pycache__/detface.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf027bccdf1a974cb2f4aae60b40ec517343ae12 Binary files /dev/null and b/hymm_sp/data_kits/face_align/__pycache__/detface.cpython-310.pyc differ diff --git a/hymm_sp/data_kits/face_align/align.py b/hymm_sp/data_kits/face_align/align.py new file mode 100644 index 0000000000000000000000000000000000000000..610c441efb41fa3e02fc0e8005f6aff42b80d333 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/hymm_sp/data_kits/face_align/detface.py b/hymm_sp/data_kits/face_align/detface.py new file mode 100644 index 0000000000000000000000000000000000000000..d04d2935b3bc298a0d5ee5e2e6021ef1d7e6db19 --- /dev/null +++ b/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 diff --git a/hymm_sp/data_kits/ffmpeg_utils.py b/hymm_sp/data_kits/ffmpeg_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1bd40878a227cb86758ea761fb9114dc64ee7013 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/hymm_sp/diffusion/__init__.py b/hymm_sp/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd7680ce5ecbb6b4cba1d9cf93dc8b14faca1d5e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/hymm_sp/diffusion/__pycache__/__init__.cpython-310.pyc b/hymm_sp/diffusion/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e88a3c5dae1493897c40953e2c7b72f1d71763a Binary files /dev/null and b/hymm_sp/diffusion/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/diffusion/pipelines/__init__.py b/hymm_sp/diffusion/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94b060dbc1c93c351ab4f78feb0b9ae508dc33d2 --- /dev/null +++ b/hymm_sp/diffusion/pipelines/__init__.py @@ -0,0 +1 @@ +from .pipeline_hunyuan_video_audio import HunyuanVideoAudioPipeline diff --git a/hymm_sp/diffusion/pipelines/__pycache__/__init__.cpython-310.pyc b/hymm_sp/diffusion/pipelines/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bc26024921a532d16d17ebde9ac5d1951829806 Binary files /dev/null and b/hymm_sp/diffusion/pipelines/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/diffusion/pipelines/__pycache__/pipeline_hunyuan_video_audio.cpython-310.pyc b/hymm_sp/diffusion/pipelines/__pycache__/pipeline_hunyuan_video_audio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ac441731cb049dc2c72f318b6ba2a1a0bbb7fee Binary files /dev/null and b/hymm_sp/diffusion/pipelines/__pycache__/pipeline_hunyuan_video_audio.cpython-310.pyc differ diff --git a/hymm_sp/diffusion/pipelines/pipeline_hunyuan_video_audio.py b/hymm_sp/diffusion/pipelines/pipeline_hunyuan_video_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..b9086e08156e7e8c5a110b2473d73e805b42d960 --- /dev/null +++ b/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) diff --git a/hymm_sp/diffusion/schedulers/__init__.py b/hymm_sp/diffusion/schedulers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6afd057baf947ecf3777405c7ecf60a68277365c --- /dev/null +++ b/hymm_sp/diffusion/schedulers/__init__.py @@ -0,0 +1 @@ +from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler \ No newline at end of file diff --git a/hymm_sp/diffusion/schedulers/__pycache__/__init__.cpython-310.pyc b/hymm_sp/diffusion/schedulers/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a667aaaef47e0754e5f7745214995c79d4ac7a86 Binary files /dev/null and b/hymm_sp/diffusion/schedulers/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/diffusion/schedulers/__pycache__/scheduling_flow_match_discrete.cpython-310.pyc b/hymm_sp/diffusion/schedulers/__pycache__/scheduling_flow_match_discrete.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ed7c5162afe9af8c29ac51f34c4c555a6e5285d Binary files /dev/null and b/hymm_sp/diffusion/schedulers/__pycache__/scheduling_flow_match_discrete.cpython-310.pyc differ diff --git a/hymm_sp/diffusion/schedulers/scheduling_flow_match_discrete.py b/hymm_sp/diffusion/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..194d926493b5e3800282dff1a773a75511e8c6f4 --- /dev/null +++ b/hymm_sp/diffusion/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,240 @@ +# Copyright 2024 Stability AI, Katherine Crowson and 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 +# +# ============================================================================== + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.schedulers.scheduling_utils import SchedulerMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + n_tokens: Optional[int] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None, + n_tokens: int = None): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.float() * dt + else: + raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}") + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/hymm_sp/helpers.py b/hymm_sp/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..3e29c5aa2fad008d2b61598d4e54afa5d47e4fdb --- /dev/null +++ b/hymm_sp/helpers.py @@ -0,0 +1,103 @@ +import torch +from typing import Union, List +from hymm_sp.modules.posemb_layers import get_1d_rotary_pos_embed, get_meshgrid_nd + +from itertools import repeat +import collections.abc + + +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + return parse + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + +def get_rope_freq_from_size(latents_size, ndim, target_ndim, args, + rope_theta_rescale_factor: Union[float, List[float]]=1.0, + rope_interpolation_factor: Union[float, List[float]]=1.0, + concat_dict={}): + + if isinstance(args.patch_size, int): + assert all(s % args.patch_size == 0 for s in latents_size), \ + f"Latent size(last {ndim} dimensions) should be divisible by patch size({args.patch_size}), " \ + f"but got {latents_size}." + rope_sizes = [s // args.patch_size for s in latents_size] + elif isinstance(args.patch_size, list): + assert all(s % args.patch_size[idx] == 0 for idx, s in enumerate(latents_size)), \ + f"Latent size(last {ndim} dimensions) should be divisible by patch size({args.patch_size}), " \ + f"but got {latents_size}." + rope_sizes = [s // args.patch_size[idx] for idx, s in enumerate(latents_size)] + + if len(rope_sizes) != target_ndim: + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis + head_dim = args.hidden_size // args.num_heads + rope_dim_list = args.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert sum(rope_dim_list) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed_new(rope_dim_list, + rope_sizes, + theta=args.rope_theta, + use_real=True, + theta_rescale_factor=rope_theta_rescale_factor, + interpolation_factor=rope_interpolation_factor, + concat_dict=concat_dict) + return freqs_cos, freqs_sin + +def get_nd_rotary_pos_embed_new(rope_dim_list, start, *args, theta=10000., use_real=False, + theta_rescale_factor: Union[float, List[float]]=1.0, + interpolation_factor: Union[float, List[float]]=1.0, + concat_dict={} + ): + + grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + if len(concat_dict)<1: + pass + else: + if concat_dict['mode']=='timecat': + bias = grid[:,:1].clone() + bias[0] = concat_dict['bias']*torch.ones_like(bias[0]) + grid = torch.cat([bias, grid], dim=1) + + elif concat_dict['mode']=='timecat-w': + bias = grid[:,:1].clone() + bias[0] = concat_dict['bias']*torch.ones_like(bias[0]) + bias[2] += start[-1] ## ref https://github.com/Yuanshi9815/OminiControl/blob/main/src/generate.py#L178 + grid = torch.cat([bias, grid], dim=1) + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len(rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len(rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed(rope_dim_list[i], grid[i].reshape(-1), theta, use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i]) # 2 x [WHD, rope_dim_list[i]] + + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb \ No newline at end of file diff --git a/hymm_sp/inference.py b/hymm_sp/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..133d91d281c3130fa27cabffaf20887000a7a373 --- /dev/null +++ b/hymm_sp/inference.py @@ -0,0 +1,178 @@ +import torch +from pathlib import Path +from loguru import logger +from hymm_sp.constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE +from hymm_sp.vae import load_vae +from hymm_sp.modules import load_model +from hymm_sp.text_encoder import TextEncoder +import torch.distributed +from hymm_sp.modules.parallel_states import ( + nccl_info, +) +from hymm_sp.modules.fp8_optimization import convert_fp8_linear + + +class Inference(object): + def __init__(self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + cpu_offload=False, + device=None, + logger=None): + self.vae = vae + self.vae_kwargs = vae_kwargs + + self.text_encoder = text_encoder + self.text_encoder_2 = text_encoder_2 + + self.model = model + self.pipeline = pipeline + self.cpu_offload = cpu_offload + + self.args = args + self.device = device if device is not None else "cuda" if torch.cuda.is_available() else "cpu" + if nccl_info.sp_size > 1: + self.device = torch.device(f"cuda:{torch.distributed.get_rank()}") + + self.logger = logger + + @classmethod + def from_pretrained(cls, + pretrained_model_path, + args, + device=None, + **kwargs): + """ + Initialize the Inference pipeline. + + Args: + pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints. + device (int): The device for inference. Default is 0. + logger (logging.Logger): The logger for the inference pipeline. Default is None. + """ + # ======================================================================== + logger.info(f"Got text-to-video model root path: {pretrained_model_path}") + + # ======================== Get the args path ============================= + + # Set device and disable gradient + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + torch.set_grad_enabled(False) + logger.info("Building model...") + factor_kwargs = {'device': 'cpu' if args.cpu_offload else device, 'dtype': PRECISION_TO_TYPE[args.precision]} + in_channels = args.latent_channels + out_channels = args.latent_channels + print("="*25, f"build model", "="*25) + model = load_model( + args, + in_channels=in_channels, + out_channels=out_channels, + factor_kwargs=factor_kwargs + ) + if args.use_fp8: + convert_fp8_linear(model, pretrained_model_path, original_dtype=PRECISION_TO_TYPE[args.precision]) + if args.cpu_offload: + print(f'='*20, f'load transformer to cpu') + model = model.to('cpu') + torch.cuda.empty_cache() + else: + model = model.to(device) + model = Inference.load_state_dict(args, model, pretrained_model_path) + model.eval() + + # ============================= Build extra models ======================== + # VAE + print("="*25, f"load vae", "="*25) + vae, _, s_ratio, t_ratio = load_vae(args.vae, args.vae_precision, logger=logger, device='cpu' if args.cpu_offload else device) + vae_kwargs = {'s_ratio': s_ratio, 't_ratio': t_ratio} + + # Text encoder + if args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0) + else: + crop_start = 0 + max_length = args.text_len + crop_start + + # prompt_template_video + prompt_template_video = PROMPT_TEMPLATE[args.prompt_template_video] if args.prompt_template_video is not None else None + print("="*25, f"load llava", "="*25) + text_encoder = TextEncoder(text_encoder_type = args.text_encoder, + max_length = max_length, + text_encoder_precision = args.text_encoder_precision, + tokenizer_type = args.tokenizer, + use_attention_mask = args.use_attention_mask, + prompt_template_video = prompt_template_video, + hidden_state_skip_layer = args.hidden_state_skip_layer, + apply_final_norm = args.apply_final_norm, + reproduce = args.reproduce, + logger = logger, + device = 'cpu' if args.cpu_offload else device , + ) + text_encoder_2 = None + if args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder(text_encoder_type=args.text_encoder_2, + max_length=args.text_len_2, + text_encoder_precision=args.text_encoder_precision_2, + tokenizer_type=args.tokenizer_2, + use_attention_mask=args.use_attention_mask, + reproduce=args.reproduce, + logger=logger, + device='cpu' if args.cpu_offload else device , # if not args.use_cpu_offload else 'cpu' + ) + + return cls(args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + model=model, + text_encoder_2=text_encoder_2, + device=device, + logger=logger) + + @staticmethod + def load_state_dict(args, model, ckpt_path): + load_key = args.load_key + ckpt_path = Path(ckpt_path) + if ckpt_path.is_dir(): + ckpt_path = next(ckpt_path.glob("*_model_states.pt")) + state_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage) + if load_key in state_dict: + state_dict = state_dict[load_key] + elif load_key == ".": + pass + else: + raise KeyError(f"Key '{load_key}' not found in the checkpoint. Existed keys: {state_dict.keys()}") + model.load_state_dict(state_dict, strict=False) + return model + + def get_exp_dir_and_ckpt_id(self): + if self.ckpt is None: + raise ValueError("The checkpoint path is not provided.") + + ckpt = Path(self.ckpt) + if ckpt.parents[1].name == "checkpoints": + # It should be a standard checkpoint path. We use the parent directory as the default save directory. + exp_dir = ckpt.parents[2] + else: + raise ValueError(f"We cannot infer the experiment directory from the checkpoint path: {ckpt}. " + f"It seems that the checkpoint path is not standard. Please explicitly provide the " + f"save path by --save-path.") + return exp_dir, ckpt.parent.name + + @staticmethod + def parse_size(size): + if isinstance(size, int): + size = [size] + if not isinstance(size, (list, tuple)): + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + if len(size) == 1: + size = [size[0], size[0]] + if len(size) != 2: + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + return size diff --git a/hymm_sp/modules/__init__.py b/hymm_sp/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9afb3ec4a0ff741c449f0f5395607f1d389a8898 --- /dev/null +++ b/hymm_sp/modules/__init__.py @@ -0,0 +1,11 @@ +from .models_audio import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG + +def load_model(args, in_channels, out_channels, factor_kwargs): + model = HYVideoDiffusionTransformer( + args, + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + return model diff --git a/hymm_sp/modules/__pycache__/__init__.cpython-310.pyc b/hymm_sp/modules/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9a48e7e8854f0d8ce96e4290f145be21bb77b08 Binary files /dev/null and b/hymm_sp/modules/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/activation_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/activation_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99476634d6ab8c76c0e79fb8c415a3f29761884c Binary files /dev/null and b/hymm_sp/modules/__pycache__/activation_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/attn_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/attn_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af0890961025d79bcfec925fc232c489b2244f9f Binary files /dev/null and b/hymm_sp/modules/__pycache__/attn_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/audio_adapters.cpython-310.pyc b/hymm_sp/modules/__pycache__/audio_adapters.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a87722e5686e14d1829539014b1e7c4893bb47da Binary files /dev/null and b/hymm_sp/modules/__pycache__/audio_adapters.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/embed_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/embed_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bca5dd2426b1959a6ce2ee4a04ad9a6b65901f29 Binary files /dev/null and b/hymm_sp/modules/__pycache__/embed_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/fp8_optimization.cpython-310.pyc b/hymm_sp/modules/__pycache__/fp8_optimization.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1db5858f98759ec2147fcc932c00b75c9688f5ca Binary files /dev/null and b/hymm_sp/modules/__pycache__/fp8_optimization.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/mlp_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/mlp_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f671be11cbc484fa4e042d7ec15e8ac2ad981178 Binary files /dev/null and b/hymm_sp/modules/__pycache__/mlp_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/models_audio.cpython-310.pyc b/hymm_sp/modules/__pycache__/models_audio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4df97227b143b7db6ef8e6c33be195b2d947c78 Binary files /dev/null and b/hymm_sp/modules/__pycache__/models_audio.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/modulate_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/modulate_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab2cc31965d11618f90ad6d5b9750ef2a9e570e5 Binary files /dev/null and b/hymm_sp/modules/__pycache__/modulate_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/norm_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/norm_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..981324d8c751bfaf1c8691dc53bce4bfab39bc8a Binary files /dev/null and b/hymm_sp/modules/__pycache__/norm_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/parallel_states.cpython-310.pyc b/hymm_sp/modules/__pycache__/parallel_states.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1116296a6bd75ea6e8acfdcba94dbc562bdb05ad Binary files /dev/null and b/hymm_sp/modules/__pycache__/parallel_states.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/posemb_layers.cpython-310.pyc b/hymm_sp/modules/__pycache__/posemb_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..133696dba2d96e6b33c0c9adb447bb290ebba69b Binary files /dev/null and b/hymm_sp/modules/__pycache__/posemb_layers.cpython-310.pyc differ diff --git a/hymm_sp/modules/__pycache__/token_refiner.cpython-310.pyc b/hymm_sp/modules/__pycache__/token_refiner.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efa0256838f28ec95d8d973e066cf6917d2e51c3 Binary files /dev/null and b/hymm_sp/modules/__pycache__/token_refiner.cpython-310.pyc differ diff --git a/hymm_sp/modules/activation_layers.py b/hymm_sp/modules/activation_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd0e145ca844b5321df35dabe700e335487a3a35 --- /dev/null +++ b/hymm_sp/modules/activation_layers.py @@ -0,0 +1,23 @@ +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") \ No newline at end of file diff --git a/hymm_sp/modules/attn_layers.py b/hymm_sp/modules/attn_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..d51e2ef24a248bf11554862ec910de848e0cb806 --- /dev/null +++ b/hymm_sp/modules/attn_layers.py @@ -0,0 +1,429 @@ +import importlib.metadata +import math +from typing import Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +try: + from flash_attn import flash_attn_qkvpacked_func, flash_attn_kvpacked_func, flash_attn_varlen_kvpacked_func + from flash_attn.bert_padding import index_first_axis +except ImportError: + flash_attn_qkvpacked_func, flash_attn_kvpacked_func, flash_attn_varlen_kvpacked_func = None, None, None + index_first_axis = None +from packaging import version +from transformers.utils.import_utils import _is_package_available + +from .norm_layers import get_norm_layer + + +def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x: torch.Tensor, head_first=False): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}' + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis[0].shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}' + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}' + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis.shape == (x.shape[1], x.shape[-1]), f'freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}' + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +class BasicAttentionLayer(nn.Module): + def __init__(self, attn_mode='flash', deterministic=False): + super().__init__() + self.attn_mode = attn_mode + self.deterministic = deterministic + + def set_attn_mode(self, new_mode): + self.attn_mode = new_mode + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + +MEMORY_LAYOUT = { + "self_flash": ( + lambda x: x, + lambda x: x, + ), + "cross_flash": ( + lambda x: x, + lambda x: x, + ), + "torch": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), + "vanilla": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), +} + + +# Copyed from https://github.com/huggingface/transformers/blob/b873234cb649a24865021f0d598627ce2b24d34a/src/transformers/modeling_flash_attention_utils.py#L33C1-L57C6 +def _get_unpad_data(attention_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, int]: + """ + Retrieves indexing data required to repad unpadded (ragged) tensors. + + Arguments: + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + indices (`torch.Tensor): + The indices of non-masked tokens from the flattened input sequence. + cu_seqlens (`torch.Tensor`): + The cumulative sequence lengths, used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + max_seqlen_in_batch (`int`): + Maximum sequence length in batch. + """ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +# Copyed from https://github.com/huggingface/transformers/blob/b873234cb649a24865021f0d598627ce2b24d34a/src/transformers/utils/import_utils.py#L822 +def is_flash_attn_greater_or_equal(library_version: str): + if not _is_package_available("flash_attn"): + return False + + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse(library_version) + + +def get_kv_seqlens_with_mask(attn_mask, k, v): + indices_k, cu_seqlens_k, max_seqlen_k = _get_unpad_data(attn_mask) + b, s1, a, d = k.shape + k = index_first_axis(k.reshape(b * s1, a, d), indices_k) + v = index_first_axis(v.reshape(b * s1, a, d), indices_k) + kv = torch.stack([k, v], dim=1) + return cu_seqlens_k, max_seqlen_k, kv + + +def get_q_seqlens(q): + bs, s, a, d = q.shape + cu_seqlens_q = torch.arange(0, (bs + 1) * s, step=s, dtype=torch.int32, device=q.device) + q = q.reshape(bs * s, a, d) + return cu_seqlens_q, s, q + + +def attention(q, k, v, mode, drop_rate=0, attn_mask=None, causal=False, deterministic=False, + cu_seqlens=None, max_seqlen=None, cu_seqlens_k=None, max_seqlen_k=None): + """ + Perform QKV self attention. + + Args: + q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads. + k (torch.Tensor): Key tensor with shape [b, s1, a, d] + v (torch.Tensor): Value tensor with shape [b, s1, a, d] + mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'. + drop_rate (float): Dropout rate in attention map. (default: 0) + attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla). + (default: None) + causal (bool): Whether to use causal attention. (default: False) + deterministic (bool): Whether to use deterministic attention. (default: False) + cu_seqlens (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch, + used to index into q. + max_seqlen (int): The maximum sequence length in the batch of q. + cu_seqlens_k (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch, + used to index into kv. + max_seqlen_k (int): The maximum sequence length in the batch of k and v. + + Returns: + torch.Tensor: Output tensor after self attention with shape [b, s, ad] + """ + pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode] + q = pre_attn_layout(q) + k = pre_attn_layout(k) + v = pre_attn_layout(v) + + if mode == 'torch': + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) + + elif mode == 'vanilla': + scale_factor = 1 / math.sqrt(q.size(-1)) + + b, a, s, _ = q.shape + s1 = k.size(2) + attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device) + if causal: + # Only applied to self attention + assert attn_mask is None, "Causal mask and attn_mask cannot be used together" + temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(diagonal=0) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(q.dtype) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias += attn_mask + + attn = (q @ k.transpose(-2, -1)) * scale_factor + attn += attn_bias + attn = attn.softmax(dim=-1) + attn = torch.dropout(attn, p=drop_rate, train=True) + x = attn @ v + else: + raise NotImplementedError(f'Unsupported attention mode: {mode}') + + x = post_attn_layout(x) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +class SelfAttentionLayer(BasicAttentionLayer): + def __init__(self, + dim, + num_heads, + qkv_bias=True, + qk_norm=True, + attn_drop=0, + proj_drop=0, + dtype=None, + device=None, + norm_type='layer', + attn_mode='self_flash', + deterministic=False, + ) -> None: + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__(attn_mode, deterministic) + self.dim = dim + self.num_heads = num_heads + assert self.dim % num_heads == 0, "dim must be divisible by num_heads" + self.head_dim = self.dim // num_heads + self.attn_drop = attn_drop + + # This assertion is aligned with flash attention + assert ( + self.head_dim % 8 == 0 and self.head_dim <= 128 + ), "Only support head_dim <= 128 and divisible by 8" + + self.Wqkv = nn.Linear(dim, dim * 3, bias=qkv_bias, **factory_kwargs) + + norm_layer = get_norm_layer(norm_type) + self.q_norm = ( + norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + + self.out_proj = nn.Linear(dim, dim, bias=qkv_bias, **factory_kwargs) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x, freqs_cis=None, attn_mask=None): + """ + Args: + x (torch.Tensor): (batch, seq_len, hidden_dim) (where hidden_dim = num heads * head dim) + freqs_cis (torch.Tensor, optional): (batch, hidden_dim // 2), RoPE for image + attn_mask (torch.Tensor, optional): (batch, seq_len, seq_len), mask for attention + """ + b, s, d = x.shape + + # Apply QKV projection + qkv = self.Wqkv(x) + qkv = qkv.view(b, s, 3, self.num_heads, self.head_dim) # [b, s, 3, a, d] + q, k, v = qkv.unbind(dim=2) # [b, s, a, d] + + # Apply QK-Norm if needed + q = self.q_norm(q) + k = self.k_norm(k) + + # Apply RoPE if needed + if freqs_cis is not None: + qq, kk = apply_rotary_emb(q, k, freqs_cis) + assert qq.shape == q.shape and kk.shape == k.shape, \ + f'qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}' + q, k = qq, kk + + # Apply self attention + context = attention(q, k, v, + drop_rate=self.attn_drop if self.training else 0, + attn_mask=attn_mask, + mode=self.attn_mode, + deterministic=self.deterministic, + ) + out = self.out_proj(context) + out = self.proj_drop(out) + + return out + + +class CrossAttentionLayer(BasicAttentionLayer): + def __init__(self, + qdim, + kdim, + num_heads, + qkv_bias=True, + qk_norm=True, + attn_drop=0, + proj_drop=0, + dtype=None, + device=None, + norm_type='layer', + attn_mode='cross_flash', + deterministic=False, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__(attn_mode, deterministic) + self.qdim = qdim + self.kdim = kdim + self.num_heads = num_heads + assert self.qdim % num_heads == 0, "qdim must be divisible by num_heads" + self.head_dim = self.qdim // num_heads + self.attn_drop = attn_drop + + # This assertion is aligned with flash attention + assert ( + self.head_dim % 8 == 0 and self.head_dim <= 128 + ), "Only support head_dim <= 128 and divisible by 8" + + self.q_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs) + self.kv_proj = nn.Linear(kdim, 2 * qdim, bias=qkv_bias, **factory_kwargs) + + norm_layer = get_norm_layer(norm_type) + self.q_norm = ( + norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + + self.out_proj = nn.Linear(qdim, qdim, bias=qkv_bias, **factory_kwargs) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x, y, attn_mask=None): + """ + Args: + x (torch.Tensor): (batch, seq_len, hidden_dim) (where hidden_dim = num heads * head dim) + y (torch.Tensor): (batch, seq_len1, hidden_dim1) + attn_mask (torch.Tensor): (batch, seq_len1), mask for attention + """ + b, s, d = x.shape + _, s1, d1 = y.shape + + q = self.q_proj(x).view(b, s, self.num_heads, self.head_dim) + kv = self.kv_proj(y).view(b, s1, 2, self.num_heads, self.head_dim) + k, v = kv.unbind(dim=2) + + # Apply QK-Norm if needed + q = self.q_norm(q) + k = self.k_norm(k) + + # Apply cross attention + context = attention(q, k, v, + attn_mask=attn_mask, + drop_rate=self.attn_drop if self.training else 0, + mode=self.attn_mode, + deterministic=self.deterministic, + ) + out = self.out_proj(context) + out = self.proj_drop(out) + + return out diff --git a/hymm_sp/modules/audio_adapters.py b/hymm_sp/modules/audio_adapters.py new file mode 100644 index 0000000000000000000000000000000000000000..00a6f9ed61c9faf9c1f5322a691c15de741f9692 --- /dev/null +++ b/hymm_sp/modules/audio_adapters.py @@ -0,0 +1,228 @@ +""" +This module provides the implementation of an Audio Projection Model, which is designed for +audio processing tasks. The model takes audio embeddings as input and outputs context tokens +that can be used for various downstream applications, such as audio analysis or synthesis. + +The AudioProjModel class is based on the ModelMixin class from the diffusers library, which +provides a foundation for building custom models. This implementation includes multiple linear +layers with ReLU activation functions and a LayerNorm for normalization. + +Key Features: +- Audio embedding input with flexible sequence length and block structure. +- Multiple linear layers for feature transformation. +- ReLU activation for non-linear transformation. +- LayerNorm for stabilizing and speeding up training. +- Rearrangement of input embeddings to match the model's expected input shape. +- Customizable number of blocks, channels, and context tokens for adaptability. + +The module is structured to be easily integrated into larger systems or used as a standalone +component for audio feature extraction and processing. + +Classes: +- AudioProjModel: A class representing the audio projection model with configurable parameters. + +Functions: +- (none) + +Dependencies: +- torch: For tensor operations and neural network components. +- diffusers: For the ModelMixin base class. +- einops: For tensor rearrangement operations. + +""" + +import torch +from diffusers import ModelMixin +from einops import rearrange + +import math +import torch.nn as nn +from .parallel_states import ( + initialize_sequence_parallel_state, + nccl_info, + get_sequence_parallel_state, + parallel_attention, + all_gather, + all_to_all_4D, +) + +class AudioProjNet2(ModelMixin): + """Audio Projection Model + + This class defines an audio projection model that takes audio embeddings as input + and produces context tokens as output. The model is based on the ModelMixin class + and consists of multiple linear layers and activation functions. It can be used + for various audio processing tasks. + + Attributes: + seq_len (int): The length of the audio sequence. + blocks (int): The number of blocks in the audio projection model. + channels (int): The number of channels in the audio projection model. + intermediate_dim (int): The intermediate dimension of the model. + context_tokens (int): The number of context tokens in the output. + output_dim (int): The output dimension of the context tokens. + + Methods: + __init__(self, seq_len=5, blocks=12, channels=768, intermediate_dim=512, context_tokens=32, output_dim=768): + Initializes the AudioProjModel with the given parameters. + forward(self, audio_embeds): + Defines the forward pass for the AudioProjModel. + Parameters: + audio_embeds (torch.Tensor): The input audio embeddings with shape (batch_size, video_length, blocks, channels). + Returns: + context_tokens (torch.Tensor): The output context tokens with shape (batch_size, video_length, context_tokens, output_dim). + + """ + + def __init__( + self, + seq_len=5, + blocks=12, # add a new parameter blocks + channels=768, # add a new parameter channels + intermediate_dim=512, + output_dim=768, + context_tokens=4, + ): + super().__init__() + + self.seq_len = seq_len + self.blocks = blocks + self.channels = channels + self.input_dim = ( + seq_len * blocks * channels + ) + self.intermediate_dim = intermediate_dim + self.context_tokens = context_tokens + self.output_dim = output_dim + + # define multiple linear layers + self.proj1 = nn.Linear(self.input_dim, intermediate_dim) + self.proj2 = nn.Linear(intermediate_dim, intermediate_dim) + self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim) + + self.norm = nn.LayerNorm(output_dim) + + + def forward(self, audio_embeds): + + video_length = audio_embeds.shape[1] + audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c") + batch_size, window_size, blocks, channels = audio_embeds.shape + audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels) + + audio_embeds = torch.relu(self.proj1(audio_embeds)) + audio_embeds = torch.relu(self.proj2(audio_embeds)) + + context_tokens = self.proj3(audio_embeds).reshape( + batch_size, self.context_tokens, self.output_dim + ) + context_tokens = self.norm(context_tokens) + out_all = rearrange( + context_tokens, "(bz f) m c -> bz f m c", f=video_length + ) + + return out_all + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + # (bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttentionCA(nn.Module): + def __init__(self, *, dim=3072, dim_head=1024, heads=33): + super().__init__() + self.scale = dim_head ** -0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head #* heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + import torch.nn.init as init + init.zeros_(self.to_out.weight) + if self.to_out.bias is not None: + init.zeros_(self.to_out.bias) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, t, aa, D) + latent (torch.Tensor): latent features + shape (b, t, hw, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + # print("latents shape: ", latents.shape) + # print("x shape: ", x.shape) + q = self.to_q(latents) + k, v = self.to_kv(x).chunk(2, dim=-1) + + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + # out = out.permute(0, 2, 1, 3) + return self.to_out(out) + #def forward(self, x, latents): + # """ + # Args: + # x (torch.Tensor): image features + # shape (b, t, aa, D) + # latent (torch.Tensor): latent features + # shape (b, t, hw, D) + # """ + # if get_sequence_parallel_state(): + # sp_size = nccl_info.sp_size + # sp_rank = nccl_info.rank_within_group + # print("rank:", latents.shape, sp_size, sp_rank) + # latents = torch.chunk(latents, sp_size, dim=1)[sp_rank] + + # x = self.norm1(x) + # latents = self.norm2(latents) + # # print("latents shape: ", latents.shape) + # # print("x shape: ", x.shape) + # q = self.to_q(latents) + # k, v = self.to_kv(x).chunk(2, dim=-1) + + # # print("q, k, v: ", q.shape, k.shape, v.shape) + + # # attention + # #scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + # #weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + # #weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + # #out = weight @ v + # def shrink_head(encoder_state, dim): + # local_heads = encoder_state.shape[dim] // nccl_info.sp_size + # return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + + # if get_sequence_parallel_state(): + # # batch_size, seq_len, attn_heads, head_dim + # q = all_to_all_4D(q, scatter_dim=2, gather_dim=1) # [2, 32256, 24, 128] + # k = shrink_head(k ,dim=2) + # v = shrink_head(v ,dim=2) + # qkv = torch.stack([query, key, value], dim=2) + # attn = flash_attn_no_pad(qkv, causal=False, dropout_p=0.0, softmax_scale=None) + # # out = out.permute(0, 2, 1, 3) + # #b, s, a, d = attn.shape + # #attn = attn.reshape(b, s, -1) + # + # out = self.to_out(attn) + # if get_sequence_parallel_state(): + # out = all_gather(out, dim=1) + # return out diff --git a/hymm_sp/modules/embed_layers.py b/hymm_sp/modules/embed_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..01e42180c212656fd871507035da0d63ef5b93bc --- /dev/null +++ b/hymm_sp/modules/embed_layers.py @@ -0,0 +1,135 @@ +import math +import torch +import torch.nn as nn +from hymm_sp.helpers import to_2tuple + + +class PatchEmbed(nn.Module): + """ 2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None + ): + factory_kwargs = {'dtype': dtype, 'device': device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, + **factory_kwargs) + nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1)) + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + shape = x.shape + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x, shape + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {'dtype': dtype, 'device': device} + super().__init__() + self.linear_1 = nn.Linear(in_features=in_channels, out_features=hidden_size, bias=True, **factory_kwargs) + self.act_1 = act_layer() + self.linear_2 = nn.Linear(in_features=hidden_size, out_features=hidden_size, bias=True, **factory_kwargs) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + def __init__(self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None + ): + factory_kwargs = {'dtype': dtype, 'device': device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb \ No newline at end of file diff --git a/hymm_sp/modules/fp8_optimization.py b/hymm_sp/modules/fp8_optimization.py new file mode 100644 index 0000000000000000000000000000000000000000..267abcf16f4c2a6979090cd3c10bf677e30d8f54 --- /dev/null +++ b/hymm_sp/modules/fp8_optimization.py @@ -0,0 +1,100 @@ +import os + +import torch +import torch.nn as nn +from torch.nn import functional as F + +def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1): + _bits = torch.tensor(bits) + _mantissa_bit = torch.tensor(mantissa_bit) + _sign_bits = torch.tensor(sign_bits) + M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits) + E = _bits - _sign_bits - M + bias = 2 ** (E - 1) - 1 + mantissa = 1 + for i in range(mantissa_bit - 1): + mantissa += 1 / (2 ** (i+1)) + maxval = mantissa * 2 ** (2**E - 1 - bias) + return maxval + +def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1): + """ + Default is E4M3. + """ + bits = torch.tensor(bits) + mantissa_bit = torch.tensor(mantissa_bit) + sign_bits = torch.tensor(sign_bits) + M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits) + E = bits - sign_bits - M + bias = 2 ** (E - 1) - 1 + mantissa = 1 + for i in range(mantissa_bit - 1): + mantissa += 1 / (2 ** (i+1)) + maxval = mantissa * 2 ** (2**E - 1 - bias) + minval = - maxval + minval = - maxval if sign_bits == 1 else torch.zeros_like(maxval) + input_clamp = torch.min(torch.max(x, minval), maxval) + log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0) + log_scales = 2.0 ** (log_scales - M - bias.type(x.dtype)) + # dequant + qdq_out = torch.round(input_clamp / log_scales) * log_scales + return qdq_out, log_scales + +def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1): + for i in range(len(x.shape) - 1): + scale = scale.unsqueeze(-1) + new_x = x / scale + quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits) + return quant_dequant_x, scale, log_scales + +def fp8_activation_dequant(qdq_out, scale, dtype): + qdq_out = qdq_out.type(dtype) + quant_dequant_x = qdq_out * scale.to(dtype) + return quant_dequant_x + +def fp8_linear_forward(cls, original_dtype, input): + weight_dtype = cls.weight.dtype + ##### + if cls.weight.dtype != torch.float8_e4m3fn: + maxval = get_fp_maxval() + scale = torch.max(torch.abs(cls.weight.flatten())) / maxval + linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale) + linear_weight = linear_weight.to(torch.float8_e4m3fn) + weight_dtype = linear_weight.dtype + else: + scale = cls.fp8_scale.to(cls.weight.device) + linear_weight = cls.weight + ##### + + if weight_dtype == torch.float8_e4m3fn and cls.weight.sum() != 0: + if True or len(input.shape) == 3: + cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype) + if cls.bias != None: + output = F.linear(input, cls_dequant, cls.bias) + else: + output = F.linear(input, cls_dequant) + return output + else: + return cls.original_forward(input.to(original_dtype)) + else: + return cls.original_forward(input) + +def convert_fp8_linear(module, dit_weight_path, original_dtype, params_to_keep={}): + setattr(module, "fp8_matmul_enabled", True) + + # loading fp8 mapping file + fp8_map_path = dit_weight_path.replace('.pt', '_map.pt') + if os.path.exists(fp8_map_path): + fp8_map = torch.load(fp8_map_path, map_location=lambda storage, loc: storage)['module'] + else: + raise ValueError(f"Invalid fp8_map path: {fp8_map_path}.") + + fp8_layers = [] + for key, layer in module.named_modules(): + if isinstance(layer, nn.Linear) and ('double_blocks' in key or 'single_blocks' in key): + fp8_layers.append(key) + original_forward = layer.forward + layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn)) + setattr(layer, "fp8_scale", fp8_map[key].to(dtype=original_dtype)) + setattr(layer, "original_forward", original_forward) + setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input)) \ No newline at end of file diff --git a/hymm_sp/modules/mlp_layers.py b/hymm_sp/modules/mlp_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..efaa2a6be1e8b14d4810c4c30e06e0d2a0b67641 --- /dev/null +++ b/hymm_sp/modules/mlp_layers.py @@ -0,0 +1,95 @@ +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from .modulate_layers import modulate +from hymm_sp.helpers import to_2tuple + + +class MLP(nn.Module): + """ MLP as used in Vision Transformer, MLP-Mixer and related networks + """ + def __init__(self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0., + use_conv=False, + device=None, + dtype=None + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity() + self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + if isinstance(patch_size, int): + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True, **factory_kwargs) + else: + self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * patch_size[2] * out_channels, bias=True) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs) + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/hymm_sp/modules/models_audio.py b/hymm_sp/modules/models_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..7303fd2a3949ea46ed1181dea0363ab3da59facc --- /dev/null +++ b/hymm_sp/modules/models_audio.py @@ -0,0 +1,745 @@ +from typing import List, Tuple, Optional, Union, Dict +from einops import rearrange + +import torch, os +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config +from flash_attn.flash_attn_interface import flash_attn_varlen_func + +from .activation_layers import get_activation_layer +from .norm_layers import get_norm_layer +from .embed_layers import TimestepEmbedder, PatchEmbed, TextProjection +from .attn_layers import apply_rotary_emb +from .mlp_layers import MLP, MLPEmbedder, FinalLayer +from .modulate_layers import ModulateDiT, modulate, apply_gate +from .token_refiner import SingleTokenRefiner +from .audio_adapters import AudioProjNet2, PerceiverAttentionCA + +from .parallel_states import ( + nccl_info, + get_cu_seqlens, + get_sequence_parallel_state, + parallel_attention, + all_gather, +) + +CPU_OFFLOAD = int(os.environ.get("CPU_OFFLOAD", 0)) +DISABLE_SP = int(os.environ.get("DISABLE_SP", 0)) +print(f'models: cpu_offload={CPU_OFFLOAD}, DISABLE_SP={DISABLE_SP}') + +class DoubleStreamBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_width_ratio: float, + mlp_act_type: str = 'gelu_tanh', + qk_norm: bool = True, + qk_norm_type: str = 'rms', + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + + self.deterministic = False + self.num_heads = num_heads + head_dim = hidden_size // num_heads + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT(hidden_size, factor=6, act_layer=get_activation_layer("silu"), **factory_kwargs) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.img_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs + ) + + self.txt_mod = ModulateDiT(hidden_size, factor=6, act_layer=get_activation_layer("silu"), **factory_kwargs) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.txt_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.txt_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.txt_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs + ) + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + freqs_cis: tuple = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + img_mod1_shift, img_mod1_scale, img_mod1_gate, img_mod2_shift, img_mod2_scale, img_mod2_gate = ( + self.img_mod(vec).chunk(6, dim=-1) + ) + txt_mod1_shift, txt_mod1_scale, txt_mod1_gate, txt_mod2_shift, txt_mod2_scale, txt_mod2_gate = ( + self.txt_mod(vec).chunk(6, dim=-1) + ) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Prepare image for attention. + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) + img_qkv = self.img_attn_qkv(img_modulated) + if CPU_OFFLOAD: torch.cuda.empty_cache() + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads) + # Apply QK-Norm if needed + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Apply RoPE if needed. + if freqs_cis is not None: + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert img_qq.shape == img_q.shape and img_kk.shape == img_k.shape, \ + f'img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}' + img_q, img_k = img_qq, img_kk + + # Prepare txt for attention. + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) + if CPU_OFFLOAD: torch.cuda.empty_cache() + txt_qkv = self.txt_attn_qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads) + # Apply QK-Norm if needed. + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Run actual attention. + q = torch.cat((img_q, txt_q), dim=1) + k = torch.cat((img_k, txt_k), dim=1) + v = torch.cat((img_v, txt_v), dim=1) + + # Compute attention. + if CPU_OFFLOAD or DISABLE_SP: + assert cu_seqlens_q.shape[0] == 2 * img.shape[0] + 1 + + q, k, v = [ + x.view(x.shape[0] * x.shape[1], *x.shape[2:]) + for x in [q, k, v] + ] + attn = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ) + attn = attn.view(img_k.shape[0], max_seqlen_q, -1).contiguous() + else: + attn, _ = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + ) + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Calculate the img bloks. + img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) + img = img + apply_gate(self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)), gate=img_mod2_gate) + if CPU_OFFLOAD: torch.cuda.empty_cache() + # Calculate the txt bloks. + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate(self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)), gate=txt_mod2_gate) + if CPU_OFFLOAD: torch.cuda.empty_cache() + return img, txt + + +class SingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = 'gelu_tanh', + qk_norm: bool = True, + qk_norm_type: str = 'rms', + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + + self.deterministic = False + self.hidden_size = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim**-0.5 + + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.mlp_act = get_activation_layer(mlp_act_type)() + self.modulation = ModulateDiT(hidden_size, factor=3, act_layer=get_activation_layer("silu"), **factory_kwargs) + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + ) -> torch.Tensor: + mod_shift, mod_scale, mod_gate = ( + self.modulation(vec).chunk(3, dim=-1) + ) + x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale) + if CPU_OFFLOAD: torch.cuda.empty_cache() + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Apply RoPE if needed. + if freqs_cis is not None: + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert img_qq.shape == img_q.shape and img_kk.shape == img_k.shape, \ + f'img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}' + img_q, img_k = img_qq, img_kk + q = torch.cat((img_q, txt_q), dim=1) + k = torch.cat((img_k, txt_k), dim=1) + + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Compute attention. + if CPU_OFFLOAD or DISABLE_SP: + assert cu_seqlens_q.shape[0] == 2 * x.shape[0] + 1, f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}" + # [b, s+l, a, d] -> [s+l, b, a, d] + q, k, v = [ + x.view(x.shape[0] * x.shape[1], *x.shape[2:]) + for x in [q, k, v] + ] + + attn = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ) + attn = attn.view(x.shape[0], max_seqlen_q, -1).contiguous() + else: + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + attn, _ = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + ) + if CPU_OFFLOAD: + torch.cuda.empty_cache() + tmp = torch.cat((attn, self.mlp_act(mlp)), 2) + torch.cuda.empty_cache() + output = self.linear2(tmp) + else: + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + apply_gate(output, gate=mod_gate) + + +class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): + """ + HunyuanVideo Transformer backbone + + Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline. + + Reference: + [1] Flux.1: https://github.com/black-forest-labs/flux + [2] MMDiT: http://arxiv.org/abs/2403.03206, + https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py + + """ + @register_to_config + def __init__( + self, + args, + patch_size: list = [1,2,2], + in_channels: int = 4, # Should be VAE.config.latent_channels. + out_channels: int = None, + hidden_size: int = 3072, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = 'gelu_tanh', + num_heads: int = 24, + depth_double_blocks: int = 19, + depth_single_blocks: int = 38, + rope_dim_list: List[int] = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = 'rms', + guidance_embed: bool = False, # For modulation. + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.text_projection = args.text_projection + self.text_states_dim = args.text_states_dim + self.use_attention_mask = args.use_attention_mask + self.text_states_dim_2 = args.text_states_dim_2 + + # Now we only use above configs from args. + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + + if hidden_size % num_heads != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" + ) + pe_dim = hidden_size // num_heads + if sum(rope_dim_list) != pe_dim: + raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}") + self.hidden_size = hidden_size + self.num_heads = num_heads + + # image projection + self.img_in = PatchEmbed( + self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs + ) + self.ref_in = PatchEmbed( + self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs + ) + + # text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + self.text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + self.text_states_dim, hidden_size, num_heads, depth=2, **factory_kwargs + ) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + + # time modulation + self.time_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) + + # text modulation + self.vector_in = MLPEmbedder( + self.text_states_dim_2, self.hidden_size, **factory_kwargs + ) + + # guidance modulation + self.guidance_in = TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs + ) if guidance_embed else None + + # double blocks + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs + ) + for _ in range(depth_double_blocks) + ] + ) + + # single blocks + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs + ) + for _ in range(depth_single_blocks) + ] + ) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs + ) + # -------------------- audio_proj_model -------------------- + self.audio_proj = AudioProjNet2(seq_len=10, blocks=5, channels=384, intermediate_dim=1024, output_dim=3072, context_tokens=4) + + # -------------------- motion-embeder -------------------- + self.motion_exp = TimestepEmbedder( + self.hidden_size // 4, + get_activation_layer("silu"), + **factory_kwargs + ) + self.motion_pose = TimestepEmbedder( + self.hidden_size // 4, + get_activation_layer("silu"), + **factory_kwargs + ) + + self.fps_proj = TimestepEmbedder( + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs + ) + + self.before_proj = nn.Linear(self.hidden_size, self.hidden_size) + + # -------------------- audio_insert_model -------------------- + self.double_stream_list = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19] + self.single_stream_list = [] + self.double_stream_map = {str(i): j for j, i in enumerate(self.double_stream_list)} + self.single_stream_map = {str(i): j+len(self.double_stream_list) for j, i in enumerate(self.single_stream_list)} + + self.audio_adapter_blocks = nn.ModuleList([ + PerceiverAttentionCA(dim=3072, dim_head=1024, heads=33) for _ in range(len(self.double_stream_list) + len(self.single_stream_list)) + ]) + + + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, # Should be in range(0, 1000). + ref_latents: torch.Tensor=None, + text_states: torch.Tensor = None, + text_mask: torch.Tensor = None, # Now we don't use it. + text_states_2: Optional[torch.Tensor] = None, # Text embedding for modulation. + freqs_cos: Optional[torch.Tensor] = None, + freqs_sin: Optional[torch.Tensor] = None, + guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + return_dict: bool = True, + is_cache: bool = False, + **additional_kwargs, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + out = {} + img = x + txt = text_states + bsz, _, ot, oh, ow = x.shape + tt, th, tw = ot // self.patch_size[0], oh // self.patch_size[1], ow // self.patch_size[2] + + # Prepare modulation vectors. + vec = self.time_in(t) + + motion_exp_vec = self.motion_exp(additional_kwargs["motion_exp"].view(-1)).view(x.shape[0], -1) # (b, 3072) + vec = vec + motion_exp_vec + motion_pose_vec = self.motion_pose(additional_kwargs["motion_pose"].view(-1)).view(x.shape[0], -1) # (b, 3072) + vec = vec + motion_pose_vec + fps_vec = self.fps_proj(additional_kwargs["fps"]) # (b, 3072) + vec = vec + fps_vec + audio_feature_all = self.audio_proj(additional_kwargs["audio_prompts"]) + + # text modulation + vec = vec + self.vector_in(text_states_2) + + # guidance modulation + if self.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + else: + # our timestep_embedding is merged into guidance_in(TimestepEmbedder) + vec = vec + self.guidance_in(guidance) + + if CPU_OFFLOAD: torch.cuda.empty_cache() + + # Embed image and text. + ref_latents_first = ref_latents[:, :, :1].clone() + img, shape_mask = self.img_in(img) + ref_latents,_ = self.ref_in(ref_latents) + ref_latents_first,_ = self.img_in(ref_latents_first) + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + # [b, l, h] + txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + img = self.before_proj(ref_latents) + img + + if CPU_OFFLOAD: torch.cuda.empty_cache() + + ref_length = ref_latents_first.shape[-2] # [b s c] + img = torch.cat([ref_latents_first, img], dim=-2) # t c + img_len = img.shape[1] + mask_len = img_len - ref_length + if additional_kwargs["face_mask"].shape[2] == 1: + face_mask = additional_kwargs["face_mask"].repeat(1,1,ot,1,1) # repeat if number of mask frame is 1 + else: + face_mask = additional_kwargs["face_mask"] + face_mask = torch.nn.functional.interpolate(face_mask, size=[ot, shape_mask[-2], shape_mask[-1]], mode="nearest") + face_mask = face_mask.view(-1,mask_len,1).repeat(1,1,img.shape[-1]).type_as(img) + + + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + cu_seqlens_q = get_cu_seqlens(text_mask, img_seq_len) + cu_seqlens_kv = cu_seqlens_q + max_seqlen_q = img_seq_len + txt_seq_len + max_seqlen_kv = max_seqlen_q + + if get_sequence_parallel_state(): + sp_size = nccl_info.sp_size + sp_rank = nccl_info.rank_within_group + assert img.shape[1] % sp_size == 0, f"Cannot split video sequence into ulysses SP ({sp_size}) parts evenly" + img = torch.chunk(img, sp_size, dim=1)[sp_rank] + freqs_cos = torch.chunk(freqs_cos, sp_size, dim=0)[sp_rank] + freqs_sin = torch.chunk(freqs_sin, sp_size, dim=0)[sp_rank] + + if CPU_OFFLOAD: torch.cuda.empty_cache() + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + # --------------------- Pass through DiT blocks ------------------------ + if not is_cache: + for layer_num, block in enumerate(self.double_blocks): + double_block_args = [img, txt, vec, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, freqs_cis] + img, txt = block(*double_block_args) + if CPU_OFFLOAD: torch.cuda.empty_cache() + """ insert audio feature to img """ + if layer_num in self.double_stream_list: + if get_sequence_parallel_state(): + img = all_gather(img, dim=1) + + real_img = img[:,ref_length:].clone().view(bsz, ot, -1, 3072) + real_ref_img = torch.zeros_like(img[:,:ref_length].clone()) + + audio_feature_pad = audio_feature_all[:,:1].repeat(1,3,1,1) + audio_feature_all_insert = torch.cat([audio_feature_pad, audio_feature_all], dim=1).view(bsz, ot, 16, 3072) + + double_idx = self.double_stream_map[str(layer_num)] + real_img = self.audio_adapter_blocks[double_idx](audio_feature_all_insert, real_img).view(bsz, -1, 3072) + img = img + torch.cat((real_ref_img, real_img * face_mask), dim=1) + if get_sequence_parallel_state(): + sp_size = nccl_info.sp_size + sp_rank = nccl_info.rank_within_group + assert img.shape[1] % sp_size == 0, f"Cannot split video sequence into ulysses SP ({sp_size}) parts evenly" + img = torch.chunk(img, sp_size, dim=1)[sp_rank] + + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + # Compatible with MMDiT. + if len(self.single_blocks) > 0: + for layer_num, block in enumerate(self.single_blocks): + if layer_num == (len(self.single_blocks) - 1): + # self.cache_out = x + tmp = x[:, :-txt_seq_len, ...] + if get_sequence_parallel_state(): + tmp = all_gather(tmp, dim=1) + self.cache_out = torch.cat([tmp, x[:, -txt_seq_len:, ...]], dim=1) + + single_block_args = [x, vec, txt_seq_len, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, (freqs_cos, freqs_sin)] + x = block(*single_block_args) + if CPU_OFFLOAD: torch.cuda.empty_cache() + else: + if get_sequence_parallel_state(): + sp_size = nccl_info.sp_size + sp_rank = nccl_info.rank_within_group + tmp, txt = self.cache_out[:, :-txt_seq_len], self.cache_out[:, -txt_seq_len:] + tmp = torch.chunk(tmp, sp_size, dim=1)[sp_rank] + x = torch.cat([tmp, txt], dim=1) + else: + x = self.cache_out + if len(self.single_blocks) > 0: + for layer_num, block in enumerate(self.single_blocks): + if layer_num < (len(self.single_blocks) - 1): + continue + single_block_args = [x, vec, txt_seq_len, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, (freqs_cos, freqs_sin)] + x = block(*single_block_args) + if CPU_OFFLOAD: torch.cuda.empty_cache() + + img = x[:, :-txt_seq_len, ...] + + if get_sequence_parallel_state(): + img = all_gather(img, dim=1) + img = img[:, ref_length:] + # ---------------------------- Final layer ------------------------------ + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + img = self.unpatchify(img, tt, th, tw) + + if return_dict: + out['x'] = img + return out + return img + + def unpatchify(self, x, t, h, w): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum('nthwcopq->nctohpwq', x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs + + def params_count(self): + counts = { + "double": sum([ + sum(p.numel() for p in block.img_attn_qkv.parameters()) + + sum(p.numel() for p in block.img_attn_proj.parameters()) + + sum(p.numel() for p in block.img_mlp.parameters()) + + sum(p.numel() for p in block.txt_attn_qkv.parameters()) + + sum(p.numel() for p in block.txt_attn_proj.parameters()) + + sum(p.numel() for p in block.txt_mlp.parameters()) + for block in self.double_blocks + ]), + "single": sum([ + sum(p.numel() for p in block.linear1.parameters()) + + sum(p.numel() for p in block.linear2.parameters()) + for block in self.single_blocks + ]), + "total": sum(p.numel() for p in self.parameters()), + } + counts["attn+mlp"] = counts["double"] + counts["single"] + return counts + +################################################################################# +# HunyuanVideo Configs # +################################################################################# + +HUNYUAN_VIDEO_CONFIG = { # Attn+MLP / Total + 'HYVideo-T/2': { # 9.0B / 12.5B + 'depth_double_blocks': 20, + 'depth_single_blocks': 40, + 'rope_dim_list': [16, 56, 56], + 'hidden_size': 3072, + 'num_heads': 24, + 'mlp_width_ratio': 4, + }, +} diff --git a/hymm_sp/modules/modulate_layers.py b/hymm_sp/modules/modulate_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..65f8c0d9391c84b81e378972539f2ef38fbcabe0 --- /dev/null +++ b/hymm_sp/modules/modulate_layers.py @@ -0,0 +1,76 @@ +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear( + hidden_size, factor * hidden_size, bias=True, **factory_kwargs + ) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) + elif scale is None: + return x + shift.unsqueeze(1) + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +def ckpt_wrapper(module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward \ No newline at end of file diff --git a/hymm_sp/modules/norm_layers.py b/hymm_sp/modules/norm_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..6e12e69963cedfc2e238f52def1da0d9e896148c --- /dev/null +++ b/hymm_sp/modules/norm_layers.py @@ -0,0 +1,77 @@ +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") \ No newline at end of file diff --git a/hymm_sp/modules/parallel_states.py b/hymm_sp/modules/parallel_states.py new file mode 100644 index 0000000000000000000000000000000000000000..795badafc6bd5c88f4928f7ceacda8f2a870054e --- /dev/null +++ b/hymm_sp/modules/parallel_states.py @@ -0,0 +1,366 @@ +import os +import torch +import datetime +import torch.distributed as dist +from typing import Any, Tuple +from torch import Tensor +from flash_attn.flash_attn_interface import flash_attn_varlen_func + + +class COMM_INFO: + def __init__(self): + self.group = None + self.sp_size = 1 + self.global_rank = 0 + self.rank_within_group = 0 + self.group_id = 0 + + +nccl_info = COMM_INFO() +_SEQUENCE_PARALLEL_STATE = False + + +def get_cu_seqlens(text_mask, img_len): + """Calculate cu_seqlens_q, cu_seqlens_kv using text_mask and img_len + + Args: + text_mask (torch.Tensor): the mask of text + img_len (int): the length of image + + Returns: + torch.Tensor: the calculated cu_seqlens for flash attention + """ + batch_size = text_mask.shape[0] + text_len = text_mask.sum(dim=1) + max_len = text_mask.shape[1] + img_len + + cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda") + + for i in range(batch_size): + s = text_len[i] + img_len + s1 = i * max_len + s + s2 = (i + 1) * max_len + cu_seqlens[2 * i + 1] = s1 + cu_seqlens[2 * i + 2] = s2 + + return cu_seqlens + +def initialize_sequence_parallel_state(sequence_parallel_size): + global _SEQUENCE_PARALLEL_STATE + if sequence_parallel_size > 1: + _SEQUENCE_PARALLEL_STATE = True + initialize_sequence_parallel_group(sequence_parallel_size) + else: + nccl_info.sp_size = 1 + nccl_info.global_rank = int(os.getenv("RANK", "0")) + nccl_info.rank_within_group = 0 + nccl_info.group_id = int(os.getenv("RANK", "0")) + +def get_sequence_parallel_state(): + return _SEQUENCE_PARALLEL_STATE + +def initialize_sequence_parallel_group(sequence_parallel_size): + """Initialize the sequence parallel group.""" + rank = int(os.getenv("RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + assert ( + world_size % sequence_parallel_size == 0 + ), "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format( + world_size, sequence_parallel_size) + nccl_info.sp_size = sequence_parallel_size + nccl_info.global_rank = rank + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + nccl_info.group = group + nccl_info.rank_within_group = rank - i * sequence_parallel_size + nccl_info.group_id = i + +def initialize_distributed(seed): + local_rank = int(os.getenv("RANK", 0)) + world_size = int(os.getenv("WORLD_SIZE", 1)) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl", init_method="env://", timeout=datetime.timedelta(seconds=2**31-1), world_size=world_size, rank=local_rank) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + initialize_sequence_parallel_state(world_size) + +def _all_to_all_4D(input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None) -> torch.tensor: + """ + all-to-all for QKV + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert (input.dim() == 4), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + if scatter_idx == 2 and gather_idx == 1: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = (input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs).transpose(0, 2).contiguous()) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = (input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, + hs).transpose(0, + 3).transpose(0, + 1).contiguous().reshape(seq_world_size, shard_hc, + shard_seqlen, bs, hs)) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + +class SeqAllToAll4D(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: Tensor, + scatter_idx: int, + gather_idx: int, + ) -> Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all(input_, ctx.world_size, process_group, scatter_dim, gather_dim) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + +def all_to_all( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim) + + +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim): + ctx.dim = dim + world_size = nccl_info.sp_size + group = nccl_info.group + input_size = list(input_.size()) + + ctx.input_size = input_size[dim] + + tensor_list = [torch.empty_like(input_) for _ in range(world_size)] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + world_size = nccl_info.sp_size + rank = nccl_info.rank_within_group + dim = ctx.dim + input_size = ctx.input_size + + sizes = [input_size] * world_size + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None + + +def all_gather(input_: torch.Tensor, dim: int = 1): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim) + +def parallel_attention(q, k, v, img_q_len, img_kv_len, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv,): + """ + img_q_len,img_kv_len: 32256 + text_mask: 2x256 + query: [2, 32256, 24, 128]) + encoder_query: [2, 256, 24, 128] + """ + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + rank = torch.distributed.get_rank() + if get_sequence_parallel_state(): + query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) # [2, 32256, 24, 128] + key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + sequence_length = query.size(1) # 32256 + encoder_sequence_length = encoder_query.size(1) # 256 + + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + bsz = query.shape[0] + head = query.shape[-2] + head_dim = query.shape[-1] + query, key, value = [ + x.view(x.shape[0] * x.shape[1], *x.shape[2:]) + for x in [query, key, value] + ] + hidden_states = flash_attn_varlen_func( + query, + key, + value, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ) + # B, S, 3, H, D + hidden_states = hidden_states.view(bsz, max_seqlen_q, head, head_dim).contiguous() + + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), + dim=1) + if get_sequence_parallel_state(): + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + + attn = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, _, _= attn.shape + attn = attn.reshape(b, s, -1) + return attn, None \ No newline at end of file diff --git a/hymm_sp/modules/posemb_layers.py b/hymm_sp/modules/posemb_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..7768cf6482a4b6f40f4e393d32cd9e08522b4e57 --- /dev/null +++ b/hymm_sp/modules/posemb_layers.py @@ -0,0 +1,164 @@ +import torch +from typing import Union, Tuple, List + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + +def get_nd_rotary_pos_embed(rope_dim_list, start, *args, theta=10000., use_real=False, + theta_rescale_factor: Union[float, List[float]]=1.0, + interpolation_factor: Union[float, List[float]]=1.0): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len(rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len(rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed(rope_dim_list[i], grid[i].reshape(-1), theta, use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i]) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed(dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor ** (dim / (dim - 2)) + + freqs = 1.0 / ( + theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim) + ) # [D/2] + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar( + torch.ones_like(freqs), freqs + ) # complex64 # [S, D/2] + return freqs_cis diff --git a/hymm_sp/modules/token_refiner.py b/hymm_sp/modules/token_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..aba2badf989c6f9dc9368917e57ccadb8b16664a --- /dev/null +++ b/hymm_sp/modules/token_refiner.py @@ -0,0 +1,213 @@ +from typing import Optional + +from einops import rearrange +import torch +import torch.nn as nn + +from .activation_layers import get_activation_layer +from .attn_layers import attention +from .norm_layers import get_norm_layer +from .embed_layers import TimestepEmbedder, TextProjection +from .attn_layers import attention +from .mlp_layers import MLP +from .modulate_layers import apply_gate + + +class IndividualTokenRefinerBlock(nn.Module): + def __init__( + self, + hidden_size, + num_heads, + mlp_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + self.num_heads = num_heads + head_dim = hidden_size // num_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm + else nn.Identity() + ) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs) + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, mode="torch", attn_mask=attn_mask) + + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + # FFN Layer + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + + return x + + +class IndividualTokenRefiner(nn.Module): + def __init__( + self, + hidden_size, + num_heads, + depth, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + self.blocks = nn.ModuleList([ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(depth) + ]) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ): + self_attn_mask = None + if mask is not None: + batch_size = mask.shape[0] + seq_len = mask.shape[1] + mask = mask.to(x.device) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + # batch_size x 1 x seq_len x seq_len, 1 for broadcasting of num_heads + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + # avoids self-attention weight being NaN for padding tokens + self_attn_mask[:, :, :, 0] = True + + for block in self.blocks: + x = block(x, c, self_attn_mask) + return x + + +class SingleTokenRefiner(nn.Module): + def __init__( + self, + in_channels, + hidden_size, + num_heads, + depth, + mlp_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {'device': device, 'dtype': dtype} + super().__init__() + + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs) + + act_layer = get_activation_layer(act_type) + # Build timestep embedding layer + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + # Build context embedding layer + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs) + + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + num_heads=num_heads, + depth=depth, + mlp_ratio=mlp_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ): + timestep_aware_representations = self.t_embedder(t) + + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] + context_aware_representations = ( + (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + ) + context_aware_representations = self.c_embedder(context_aware_representations) + c = timestep_aware_representations + context_aware_representations + + x = self.input_embedder(x) + + x = self.individual_token_refiner(x, c, mask) + + return x \ No newline at end of file diff --git a/hymm_sp/sample_batch.py b/hymm_sp/sample_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..fa263f57f4e7ad6bd5688cca3770fd83021a4b62 --- /dev/null +++ b/hymm_sp/sample_batch.py @@ -0,0 +1,114 @@ +import os +import torch +import numpy as np +from pathlib import Path +from loguru import logger +from einops import rearrange +import torch.distributed +from torch.utils.data.distributed import DistributedSampler +from torch.utils.data import DataLoader +from hymm_sp.config import parse_args +from hymm_sp.sample_inference_audio import HunyuanVideoSampler +from hymm_sp.data_kits.audio_dataset import VideoAudioTextLoaderVal +from hymm_sp.data_kits.data_tools import save_videos_grid +from hymm_sp.data_kits.face_align import AlignImage +from hymm_sp.modules.parallel_states import ( + initialize_distributed, + nccl_info, +) + +from transformers import WhisperModel +from transformers import AutoFeatureExtractor + +MODEL_OUTPUT_PATH = os.environ.get('MODEL_BASE') + + +def main(): + args = parse_args() + models_root_path = Path(args.ckpt) + print("*"*20) + initialize_distributed(args.seed) + if not models_root_path.exists(): + raise ValueError(f"`models_root` not exists: {models_root_path}") + print("+"*20) + # Create save folder to save the samples + save_path = args.save_path + if not os.path.exists(args.save_path): + os.makedirs(save_path, exist_ok=True) + + # Load models + rank = 0 + vae_dtype = torch.float16 + device = torch.device("cuda") + if nccl_info.sp_size > 1: + device = torch.device(f"cuda:{torch.distributed.get_rank()}") + rank = torch.distributed.get_rank() + + hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(args.ckpt, args=args, device=device) + # Get the updated args + args = hunyuan_video_sampler.args + + 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) + + feature_extractor = AutoFeatureExtractor.from_pretrained(f"{MODEL_OUTPUT_PATH}/ckpts/whisper-tiny/") + + kwargs = { + "text_encoder": hunyuan_video_sampler.text_encoder, + "text_encoder_2": hunyuan_video_sampler.text_encoder_2, + "feature_extractor": feature_extractor, + } + video_dataset = VideoAudioTextLoaderVal( + image_size=args.image_size, + meta_file=args.input, + **kwargs, + ) + + sampler = DistributedSampler(video_dataset, num_replicas=1, rank=0, shuffle=False, drop_last=False) + json_loader = DataLoader(video_dataset, batch_size=1, shuffle=False, sampler=sampler, drop_last=False) + + for batch_index, batch in enumerate(json_loader, start=1): + + fps = batch["fps"] + videoid = batch['videoid'][0] + audio_path = str(batch["audio_path"][0]) + save_path = args.save_path + output_path = f"{save_path}/{videoid}.mp4" + output_audio_path = f"{save_path}/{videoid}_audio.mp4" + + samples = hunyuan_video_sampler.predict(args, batch, wav2vec, feature_extractor, align_instance) + + sample = samples['samples'][0].unsqueeze(0) # denoised latent, (bs, 16, t//4, h//8, w//8) + sample = sample[:, :, :batch["audio_len"][0]] + + video = rearrange(sample[0], "c f h w -> f h w c") + video = (video * 255.).data.cpu().numpy().astype(np.uint8) # （f h w c) + + torch.cuda.empty_cache() + + final_frames = [] + for frame in video: + final_frames.append(frame) + final_frames = np.stack(final_frames, axis=0) + + if rank == 0: + from hymm_sp.data_kits.ffmpeg_utils import save_video + save_video(final_frames, output_path, n_rows=len(final_frames), fps=fps.item()) + os.system(f"ffmpeg -i '{output_path}' -i '{audio_path}' -shortest '{output_audio_path}' -y -loglevel quiet; rm '{output_path}'") + + + + +if __name__ == "__main__": + main() + + + + + + + diff --git a/hymm_sp/sample_gpu_poor.py b/hymm_sp/sample_gpu_poor.py new file mode 100644 index 0000000000000000000000000000000000000000..cdfb719f18b250e11d2217b905bbb9444d3e4073 --- /dev/null +++ b/hymm_sp/sample_gpu_poor.py @@ -0,0 +1,115 @@ +import os +import numpy as np +from pathlib import Path +from loguru import logger +import torch +from einops import rearrange +import torch.distributed +from torch.utils.data.distributed import DistributedSampler +from torch.utils.data import DataLoader +from hymm_sp.config import parse_args +from hymm_sp.sample_inference_audio import HunyuanVideoSampler +from hymm_sp.data_kits.audio_dataset import VideoAudioTextLoaderVal +from hymm_sp.data_kits.face_align import AlignImage + +from transformers import WhisperModel +from transformers import AutoFeatureExtractor + +MODEL_OUTPUT_PATH = os.environ.get('MODEL_BASE') + + +def main(): + args = parse_args() + models_root_path = Path(args.ckpt) + + if not models_root_path.exists(): + raise ValueError(f"`models_root` not exists: {models_root_path}") + + # Create save folder to save the samples + save_path = args.save_path if args.save_path_suffix=="" else f'{args.save_path}_{args.save_path_suffix}' + if not os.path.exists(args.save_path): + os.makedirs(save_path, exist_ok=True) + + # Load models + rank = 0 + vae_dtype = torch.float16 + device = torch.device("cuda") + + hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(args.ckpt, args=args, device=device) + # Get the updated args + args = hunyuan_video_sampler.args + if args.cpu_offload: + from diffusers.hooks import apply_group_offloading + onload_device = torch.device("cuda") + apply_group_offloading(hunyuan_video_sampler.pipeline.transformer, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=1) + + 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) + + feature_extractor = AutoFeatureExtractor.from_pretrained(f"{MODEL_OUTPUT_PATH}/ckpts/whisper-tiny/") + + kwargs = { + "text_encoder": hunyuan_video_sampler.text_encoder, + "text_encoder_2": hunyuan_video_sampler.text_encoder_2, + "feature_extractor": feature_extractor, + } + video_dataset = VideoAudioTextLoaderVal( + image_size=args.image_size, + meta_file=args.input, + **kwargs, + ) + + sampler = DistributedSampler(video_dataset, num_replicas=1, rank=0, shuffle=False, drop_last=False) + json_loader = DataLoader(video_dataset, batch_size=1, shuffle=False, sampler=sampler, drop_last=False) + + for batch_index, batch in enumerate(json_loader, start=1): + + fps = batch["fps"] + videoid = batch['videoid'][0] + audio_path = str(batch["audio_path"][0]) + save_path = args.save_path + output_path = f"{save_path}/{videoid}.mp4" + output_audio_path = f"{save_path}/{videoid}_audio.mp4" + + if args.infer_min: + batch["audio_len"][0] = 129 + + samples = hunyuan_video_sampler.predict(args, batch, wav2vec, feature_extractor, align_instance) + + sample = samples['samples'][0].unsqueeze(0) # denoised latent, (bs, 16, t//4, h//8, w//8) + sample = sample[:, :, :batch["audio_len"][0]] + + video = rearrange(sample[0], "c f h w -> f h w c") + video = (video * 255.).data.cpu().numpy().astype(np.uint8) # （f h w c) + + torch.cuda.empty_cache() + + final_frames = [] + for frame in video: + final_frames.append(frame) + final_frames = np.stack(final_frames, axis=0) + + if rank == 0: + from hymm_sp.data_kits.ffmpeg_utils import save_video + save_video(final_frames, output_path, n_rows=len(final_frames), fps=fps.item()) + os.system(f"ffmpeg -i '{output_path}' -i '{audio_path}' -shortest '{output_audio_path}' -y -loglevel quiet; rm '{output_path}'") + + + + + + + +if __name__ == "__main__": + main() + + + + + + + diff --git a/hymm_sp/sample_inference_audio.py b/hymm_sp/sample_inference_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..769836a28ce311ead5080b3badf455e1ff7fa0ba --- /dev/null +++ b/hymm_sp/sample_inference_audio.py @@ -0,0 +1,230 @@ +import math +import time +import torch +import random +from loguru import logger +from einops import rearrange +from hymm_sp.diffusion import load_diffusion_pipeline +from hymm_sp.helpers import get_nd_rotary_pos_embed_new +from hymm_sp.inference import Inference +from hymm_sp.diffusion.schedulers import FlowMatchDiscreteScheduler +from hymm_sp.data_kits.audio_preprocessor import encode_audio, get_facemask + +def align_to(value, alignment): + return int(math.ceil(value / alignment) * alignment) + +class HunyuanVideoSampler(Inference): + def __init__(self, args, vae, vae_kwargs, text_encoder, model, text_encoder_2=None, pipeline=None, + device=0, logger=None): + super().__init__(args, vae, vae_kwargs, text_encoder, model, text_encoder_2=text_encoder_2, + pipeline=pipeline, device=device, logger=logger) + + self.args = args + self.pipeline = load_diffusion_pipeline( + args, 0, self.vae, self.text_encoder, self.text_encoder_2, self.model, + device=self.device) + print('load hunyuan model successful... ') + + def get_rotary_pos_embed(self, video_length, height, width, concat_dict={}): + target_ndim = 3 + ndim = 5 - 2 + if '884' in self.args.vae: + latents_size = [(video_length-1)//4+1 , height//8, width//8] + else: + latents_size = [video_length , height//8, width//8] + + if isinstance(self.model.patch_size, int): + assert all(s % self.model.patch_size == 0 for s in latents_size), \ + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " \ + f"but got {latents_size}." + rope_sizes = [s // self.model.patch_size for s in latents_size] + elif isinstance(self.model.patch_size, list): + assert all(s % self.model.patch_size[idx] == 0 for idx, s in enumerate(latents_size)), \ + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " \ + f"but got {latents_size}." + rope_sizes = [s // self.model.patch_size[idx] for idx, s in enumerate(latents_size)] + + if len(rope_sizes) != target_ndim: + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis + head_dim = self.model.hidden_size // self.model.num_heads + rope_dim_list = self.model.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert sum(rope_dim_list) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed_new(rope_dim_list, + rope_sizes, + theta=self.args.rope_theta, + use_real=True, + theta_rescale_factor=1, + concat_dict=concat_dict) + return freqs_cos, freqs_sin + + @torch.no_grad() + def predict(self, + args, batch, wav2vec, feature_extractor, align_instance, + **kwargs): + """ + Predict the image from the given text. + + Args: + prompt (str or List[str]): The input text. + kwargs: + size (int): The (height, width) of the output image/video. Default is (256, 256). + video_length (int): The frame number of the output video. Default is 1. + seed (int or List[str]): The random seed for the generation. Default is a random integer. + negative_prompt (str or List[str]): The negative text prompt. Default is an empty string. + infer_steps (int): The number of inference steps. Default is 100. + guidance_scale (float): The guidance scale for the generation. Default is 6.0. + num_videos_per_prompt (int): The number of videos per prompt. Default is 1. + verbose (int): 0 for no log, 1 for all log, 2 for fewer log. Default is 1. + output_type (str): The output type of the image, can be one of `pil`, `np`, `pt`, `latent`. + Default is 'pil'. + """ + + out_dict = dict() + + prompt = batch['text_prompt'][0] + image_path = str(batch["image_path"][0]) + audio_path = str(batch["audio_path"][0]) + neg_prompt = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion, blurring, Lens changes" + # videoid = batch['videoid'][0] + fps = batch["fps"].to(self.device) + audio_prompts = batch["audio_prompts"].to(self.device) + weight_dtype = audio_prompts.dtype + + audio_prompts = [encode_audio(wav2vec, audio_feat.to(dtype=wav2vec.dtype), fps.item(), num_frames=batch["audio_len"][0]) for audio_feat in audio_prompts] + audio_prompts = torch.cat(audio_prompts, dim=0).to(device=self.device, dtype=weight_dtype) + if audio_prompts.shape[1] <= 129: + audio_prompts = torch.cat([audio_prompts, torch.zeros_like(audio_prompts[:, :1]).repeat(1,129-audio_prompts.shape[1], 1, 1, 1)], dim=1) + else: + audio_prompts = torch.cat([audio_prompts, torch.zeros_like(audio_prompts[:, :1]).repeat(1, 5, 1, 1, 1)], dim=1) + + wav2vec.to("cpu") + torch.cuda.empty_cache() + + uncond_audio_prompts = torch.zeros_like(audio_prompts[:,:129]) + motion_exp = batch["motion_bucket_id_exps"].to(self.device) + motion_pose = batch["motion_bucket_id_heads"].to(self.device) + + pixel_value_ref = batch['pixel_value_ref'].to(self.device) # (b f c h w) 取值范围[0,255] + face_masks = get_facemask(pixel_value_ref.clone(), align_instance, area=3.0) + + pixel_value_ref = pixel_value_ref.clone().repeat(1,129,1,1,1) + uncond_pixel_value_ref = torch.zeros_like(pixel_value_ref) + pixel_value_ref = pixel_value_ref / 127.5 - 1. + uncond_pixel_value_ref = uncond_pixel_value_ref * 2 - 1 + + pixel_value_ref_for_vae = rearrange(pixel_value_ref, "b f c h w -> b c f h w") + uncond_uncond_pixel_value_ref = rearrange(uncond_pixel_value_ref, "b f c h w -> b c f h w") + + pixel_value_llava = batch["pixel_value_ref_llava"].to(self.device) + pixel_value_llava = rearrange(pixel_value_llava, "b f c h w -> (b f) c h w") + uncond_pixel_value_llava = pixel_value_llava.clone() + + # ========== Encode reference latents ========== + vae_dtype = self.vae.dtype + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_dtype != torch.float32): + + if args.cpu_offload: + self.vae.to('cuda') + + self.vae.enable_tiling() + ref_latents = self.vae.encode(pixel_value_ref_for_vae.clone()).latent_dist.sample() + uncond_ref_latents = self.vae.encode(uncond_uncond_pixel_value_ref).latent_dist.sample() + self.vae.disable_tiling() + if hasattr(self.vae.config, 'shift_factor') and self.vae.config.shift_factor: + ref_latents.sub_(self.vae.config.shift_factor).mul_(self.vae.config.scaling_factor) + uncond_ref_latents.sub_(self.vae.config.shift_factor).mul_(self.vae.config.scaling_factor) + else: + ref_latents.mul_(self.vae.config.scaling_factor) + uncond_ref_latents.mul_(self.vae.config.scaling_factor) + + if args.cpu_offload: + self.vae.to('cpu') + torch.cuda.empty_cache() + + face_masks = torch.nn.functional.interpolate(face_masks.float().squeeze(2), + (ref_latents.shape[-2], + ref_latents.shape[-1]), + mode="bilinear").unsqueeze(2).to(dtype=ref_latents.dtype) + + + size = (batch['pixel_value_ref'].shape[-2], batch['pixel_value_ref'].shape[-1]) + target_length = 129 + target_height = align_to(size[0], 16) + target_width = align_to(size[1], 16) + concat_dict = {'mode': 'timecat', 'bias': -1} + # concat_dict = {} + freqs_cos, freqs_sin = self.get_rotary_pos_embed( + target_length, + target_height, + target_width, + concat_dict) + n_tokens = freqs_cos.shape[0] + + generator = torch.Generator(device=self.device).manual_seed(args.seed) + + debug_str = f""" + prompt: {prompt} + image_path: {image_path} + audio_path: {audio_path} + negative_prompt: {neg_prompt} + seed: {args.seed} + fps: {fps.item()} + infer_steps: {args.infer_steps} + target_height: {target_height} + target_width: {target_width} + target_length: {target_length} + guidance_scale: {args.cfg_scale} + """ + self.logger.info(debug_str) + pipeline_kwargs = { + "cpu_offload": args.cpu_offload + } + start_time = time.time() + samples = self.pipeline(prompt=prompt, + height=target_height, + width=target_width, + frame=target_length, + num_inference_steps=args.infer_steps, + guidance_scale=args.cfg_scale, # cfg scale + + negative_prompt=neg_prompt, + num_images_per_prompt=args.num_images, + generator=generator, + prompt_embeds=None, + + ref_latents=ref_latents, # [1, 16, 1, h//8, w//8] + uncond_ref_latents=uncond_ref_latents, + pixel_value_llava=pixel_value_llava, # [1, 3, 336, 336] + uncond_pixel_value_llava=uncond_pixel_value_llava, + face_masks=face_masks, # [b f h w] + audio_prompts=audio_prompts, + uncond_audio_prompts=uncond_audio_prompts, + motion_exp=motion_exp, + motion_pose=motion_pose, + fps=fps, + + num_videos_per_prompt=1, + attention_mask=None, + negative_prompt_embeds=None, + negative_attention_mask=None, + output_type="pil", + freqs_cis=(freqs_cos, freqs_sin), + n_tokens=n_tokens, + data_type='video', + is_progress_bar=True, + vae_ver=self.args.vae, + enable_tiling=self.args.vae_tiling, + **pipeline_kwargs + )[0] + if samples is None: + return None + out_dict['samples'] = samples + gen_time = time.time() - start_time + logger.info(f"Success, time: {gen_time}") + + wav2vec.to(self.device) + + return out_dict + diff --git a/hymm_sp/text_encoder/__init__.py b/hymm_sp/text_encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..847d239f4373885f72a7c11fa7a1d0065a6cecec --- /dev/null +++ b/hymm_sp/text_encoder/__init__.py @@ -0,0 +1,294 @@ +from dataclasses import dataclass +from typing import Optional, Tuple +from copy import deepcopy + +import torch, os +import torch.nn as nn +from transformers import ( + CLIPTextModel, CLIPTokenizer, LlavaForConditionalGeneration, + LlamaTokenizerFast +) +from transformers.utils import ModelOutput +from ..constants import TEXT_ENCODER_PATH, TOKENIZER_PATH, PRECISION_TO_TYPE + +CPU_OFFLOAD = int(os.environ.get("CPU_OFFLOAD", 0)) +print(f'text_encoder: cpu_offload={CPU_OFFLOAD}') + +def use_default(value, default): + return value if value is not None else default + +def load_text_encoder(text_encoder_type, + text_encoder_precision=None, + text_encoder_path=None, + logger=None, + device=None + ): + if text_encoder_path is None: + text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] + if logger is not None: + logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}") + + if text_encoder_type == "clipL": + text_encoder = CLIPTextModel.from_pretrained(text_encoder_path) + text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm + elif text_encoder_type == "llava-llama-3-8b": + text_encoder = LlavaForConditionalGeneration.from_pretrained(text_encoder_path, low_cpu_mem_usage=True) + text_encoder.final_layer_norm = text_encoder.language_model.model.norm + else: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + + if text_encoder_precision is not None: + text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision]) + + text_encoder.requires_grad_(False) + + if logger is not None: + logger.info(f"Text encoder to dtype: {text_encoder.dtype}") + + if device is not None: + text_encoder = text_encoder.to(device) + + return text_encoder, text_encoder_path + +def load_tokenizer(tokenizer_type, + tokenizer_path=None, + padding_side="right", + logger=None + ): + if tokenizer_path is None: + tokenizer_path = TOKENIZER_PATH[tokenizer_type] + if logger is not None: + logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") + + if tokenizer_type == "clipL": + tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77) + elif tokenizer_type == "llava-llama-3-8b": + tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_path, padding_side=padding_side) + else: + raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") + + return tokenizer, tokenizer_path + + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + + +class TextEncoder(nn.Module): + def __init__(self, + text_encoder_type: str, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + tokenizer_type: Optional[str] = None, + tokenizer_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + input_max_length: Optional[int] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + logger=None, + device=None, + ): + super().__init__() + self.text_encoder_type = text_encoder_type + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.tokenizer_type = tokenizer_type if tokenizer_type is not None else text_encoder_type + self.tokenizer_path = tokenizer_path if tokenizer_path is not None else text_encoder_path + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert use_attention_mask is True, "Attention mask is True required when training videos." + self.input_max_length = input_max_length if input_max_length is not None else max_length + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + self.logger = logger + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video, ( + f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" + ) + assert '{}' in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}" + ) + + if "clip" in text_encoder_type: + self.output_key = output_key or "pooler_output" + elif "llama" in text_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + + self.model, self.model_path = load_text_encoder( + text_encoder_type=self.text_encoder_type, + text_encoder_precision=self.precision, + text_encoder_path=self.model_path, + logger=self.logger, + device=device + ) + self.dtype = self.model.dtype + self.device = self.model.device + + self.tokenizer, self.tokenizer_path = load_tokenizer( + tokenizer_type=self.tokenizer_type, + tokenizer_path=self.tokenizer_path, + padding_side="right", + logger=self.logger + ) + + def __repr__(self): + return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + + @staticmethod + def apply_text_to_template(text, template): + """ + Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If Ture, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def text2tokens(self, text, data_type='video', name='person'): + """ + Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = 'str' + if self.use_video_template: + if data_type == 'video': + prompt_template = self.prompt_template_video["template"] + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text] + if isinstance(text[0], list): + tokenize_input_type = 'list' + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = 'list' + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + kwargs = dict(truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt") + if self.text_encoder_type == "llava-llama-3-8b": + if isinstance(text, list): + for i in range(len(text)): + text[i] = text[i] + '\nThe %s looks like' % name + elif isinstance(text, str): + text = text + '\nThe %s looks like' % name + else: + raise NotImplementedError + + if tokenize_input_type == 'str': + return self.tokenizer(text, return_length=False, return_overflowing_tokens=False, return_attention_mask=True, **kwargs, ) + elif tokenize_input_type == 'list': + return self.tokenizer.apply_chat_template(text, add_generation_prompt=True, tokenize=True, return_dict=True, **kwargs, ) + else: + raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") + + def encode(self, batch_encoding, use_attention_mask=None, output_hidden_states=False, do_sample=None, + hidden_state_skip_layer=None, return_texts=False, data_type='image'): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) + hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) + do_sample = use_default(do_sample, not self.reproduce) + if CPU_OFFLOAD: + self.model.to('cuda') + print(f'encode prompt: move text_encoder to cuda') + + attention_mask = batch_encoding["attention_mask"].to(self.model.device) if use_attention_mask else None + if 'pixel_value_llava' in batch_encoding: + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(self.model.device), + attention_mask=attention_mask, + pixel_values=batch_encoding["pixel_value_llava"].to(self.model.device), + output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None) + else: + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(self.model.device), + attention_mask=attention_mask, + output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None,) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm(last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_video_template: + if data_type == 'video': + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = attention_mask[:, crop_start:] if use_attention_mask else None + if CPU_OFFLOAD: + self.model.to('cpu') + torch.cuda.empty_cache() + print(f'encode prompt successful: move text_encoder to cpu') + if output_hidden_states: + return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward(self, text, use_attention_mask=None, output_hidden_states=False, do_sample=False, + hidden_state_skip_layer=None, return_texts=False): + batch_encoding = self.text2tokens(text) + return self.encode(batch_encoding, use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, return_texts=return_texts) diff --git a/hymm_sp/text_encoder/__pycache__/__init__.cpython-310.pyc b/hymm_sp/text_encoder/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86ad5c067a59060e696bd83cc3528abef5dbd32d Binary files /dev/null and b/hymm_sp/text_encoder/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/vae/__init__.py b/hymm_sp/vae/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8fccc2dd3c139f9510fecc51f1a89fc92fac0b6 --- /dev/null +++ b/hymm_sp/vae/__init__.py @@ -0,0 +1,51 @@ +import torch +from pathlib import Path +from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from ..constants import VAE_PATH, PRECISION_TO_TYPE + +def load_vae(vae_type, + vae_precision=None, + sample_size=None, + vae_path=None, + logger=None, + device=None + ): + if vae_path is None: + vae_path = VAE_PATH[vae_type] + vae_compress_spec, _, _ = vae_type.split("-") + length = len(vae_compress_spec) + if length == 3: + if logger is not None: + logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}") + config = AutoencoderKLCausal3D.load_config(vae_path) + if sample_size: + vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size) + else: + vae = AutoencoderKLCausal3D.from_config(config) + ckpt = torch.load(Path(vae_path) / "pytorch_model.pt", map_location=vae.device) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + # vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items()} + vae.load_state_dict(vae_ckpt) + + spatial_compression_ratio = vae.config.spatial_compression_ratio + time_compression_ratio = vae.config.time_compression_ratio + else: + raise ValueError(f"Invalid VAE model: {vae_type}. Must be 3D VAE in the format of '???-*'.") + + if vae_precision is not None: + vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision]) + + vae.requires_grad_(False) + + if logger is not None: + logger.info(f"VAE to dtype: {vae.dtype}") + + if device is not None: + vae = vae.to(device) + + # Set vae to eval mode, even though it's dropout rate is 0. + vae.eval() + + return vae, vae_path, spatial_compression_ratio, time_compression_ratio diff --git a/hymm_sp/vae/__pycache__/__init__.cpython-310.pyc b/hymm_sp/vae/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7fe33e2aef6c432562576d660b45f3dded5e717 Binary files /dev/null and b/hymm_sp/vae/__pycache__/__init__.cpython-310.pyc differ diff --git a/hymm_sp/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc b/hymm_sp/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be7c7df0804ba2924cfca0cd2a26114075d7e050 Binary files /dev/null and b/hymm_sp/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc differ diff --git a/hymm_sp/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc b/hymm_sp/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f4cb1823d3a131c56fddb148418543ededf74bc Binary files /dev/null and b/hymm_sp/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc differ diff --git a/hymm_sp/vae/__pycache__/vae.cpython-310.pyc b/hymm_sp/vae/__pycache__/vae.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c91be34842b1a77a890d6c92892a1d5992c2681 Binary files /dev/null and b/hymm_sp/vae/__pycache__/vae.cpython-310.pyc differ diff --git a/hymm_sp/vae/autoencoder_kl_causal_3d.py b/hymm_sp/vae/autoencoder_kl_causal_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..2efff5c1687c77d3f13073429e91c504b0d2807e --- /dev/null +++ b/hymm_sp/vae/autoencoder_kl_causal_3d.py @@ -0,0 +1,903 @@ +import os +import math +from typing import Dict, Optional, Tuple, Union +from dataclasses import dataclass +from torch import distributed as dist +import loguru +import torch +import torch.nn as nn +import torch.distributed + +RECOMMENDED_DTYPE = torch.float16 + +def mpi_comm(): + from mpi4py import MPI + return MPI.COMM_WORLD + +from torch import distributed as dist +def mpi_rank(): + return dist.get_rank() + +def mpi_world_size(): + return dist.get_world_size() + + +class TorchIGather: + def __init__(self): + if not torch.distributed.is_initialized(): + rank = mpi_rank() + world_size = mpi_world_size() + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['MASTER_ADDR'] = '127.0.0.1' + os.environ['MASTER_PORT'] = str(29500) + torch.cuda.set_device(rank) + torch.distributed.init_process_group('nccl') + + self.handles = [] + self.buffers = [] + + self.world_size = dist.get_world_size() + self.rank = dist.get_rank() + self.groups_ids = [] + self.group = {} + + for i in range(self.world_size): + self.groups_ids.append(tuple(range(i + 1))) + + for group in self.groups_ids: + new_group = dist.new_group(group) + self.group[group[-1]] = new_group + + + def gather(self, tensor, n_rank=None): + if n_rank is not None: + group = self.group[n_rank - 1] + else: + group = None + rank = self.rank + tensor = tensor.to(RECOMMENDED_DTYPE) + if rank == 0: + buffer = [torch.empty_like(tensor) for i in range(n_rank)] + else: + buffer = None + self.buffers.append(buffer) + handle = torch.distributed.gather(tensor, buffer, async_op=True, group=group) + self.handles.append(handle) + + def wait(self): + for handle in self.handles: + handle.wait() + + def clear(self): + self.buffers = [] + self.handles = [] + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +try: + # This diffusers is modified and packed in the mirror. + from diffusers.loaders import FromOriginalVAEMixin +except ImportError: + # Use this to be compatible with the original diffusers. + from diffusers.loaders.single_file_model import FromOriginalModelMixin as FromOriginalVAEMixin +from diffusers.utils.accelerate_utils import apply_forward_hook +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, + CROSS_ATTENTION_PROCESSORS, + Attention, + AttentionProcessor, + AttnAddedKVProcessor, + AttnProcessor, +) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from .vae import DecoderCausal3D, BaseOutput, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D + +""" +use trt need install polygraphy and onnx-graphsurgeon +python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com +""" +try: + from polygraphy.backend.trt import ( TrtRunner, EngineFromBytes) + from polygraphy.backend.common import BytesFromPath +except: + print("TrtRunner or EngineFromBytes is not available, you can not use trt engine") + +@dataclass +class DecoderOutput2(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +MODEL_OUTPUT_PATH = os.environ.get('MODEL_OUTPUT_PATH') +MODEL_BASE = os.environ.get('MODEL_BASE') + +CPU_OFFLOAD = int(os.environ.get("CPU_OFFLOAD", 0)) +DISABLE_SP = int(os.environ.get("DISABLE_SP", 0)) +print(f'vae: cpu_offload={CPU_OFFLOAD}, DISABLE_SP={DISABLE_SP}') + + +class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): + r""" + A VAE model with KL loss for encoding images into latents and decoding latent representations into images. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + in_channels (int, *optional*, defaults to 3): Number of channels in the input image. + out_channels (int, *optional*, defaults to 3): Number of channels in the output. + down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + Tuple of downsample block types. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + Tuple of upsample block types. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): + Tuple of block output channels. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. + sample_size (`int`, *optional*, defaults to `32`): Sample input size. + scaling_factor (`float`, *optional*, defaults to 0.18215): + The component-wise standard deviation of the trained latent space computed using the first batch of the + training set. This is used to scale the latent space to have unit variance when training the diffusion + model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the + diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 + / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image + Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper. + force_upcast (`bool`, *optional*, default to `True`): + If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE + can be fine-tuned / trained to a lower range without loosing too much precision in which case + `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D",), + up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D",), + block_out_channels: Tuple[int] = (64,), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + sample_tsize: int = 64, + scaling_factor: float = 0.18215, + force_upcast: float = True, + spatial_compression_ratio: int = 8, + time_compression_ratio: int = 4, + disable_causal_conv: bool = False, + mid_block_add_attention: bool = True, + mid_block_causal_attn: bool = False, + use_trt_engine: bool = False, + nccl_gather: bool = True, + engine_path: str = f"{MODEL_BASE}/HYVAE_decoder+conv_256x256xT_fp16_H20.engine", + ): + super().__init__() + + self.disable_causal_conv = disable_causal_conv + self.time_compression_ratio = time_compression_ratio + + self.encoder = EncoderCausal3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + disable_causal=disable_causal_conv, + mid_block_add_attention=mid_block_add_attention, + mid_block_causal_attn=mid_block_causal_attn, + ) + + self.decoder = DecoderCausal3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + disable_causal=disable_causal_conv, + mid_block_add_attention=mid_block_add_attention, + mid_block_causal_attn=mid_block_causal_attn, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + + + # only relevant if vae tiling is enabled + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // time_compression_ratio + + self.tile_sample_min_size = self.config.sample_size + sample_size = ( + self.config.sample_size[0] + if isinstance(self.config.sample_size, (list, tuple)) + else self.config.sample_size + ) + self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1))) + self.tile_overlap_factor = 0.25 + + use_trt_engine = False #if CPU_OFFLOAD else True + # ============= parallism related code =================== + self.parallel_decode = use_trt_engine + self.nccl_gather = nccl_gather + + # only relevant if parallel_decode is enabled + self.gather_to_rank0 = self.parallel_decode + + self.engine_path = engine_path + + self.use_trt_decoder = use_trt_engine + + @property + def igather(self): + assert self.nccl_gather and self.gather_to_rank0 + if hasattr(self, '_igather'): + return self._igather + else: + self._igather = TorchIGather() + return self._igather + + @property + def use_padding(self): + return ( + self.use_trt_decoder + # dist.gather demands all processes possess to have the same tile shape. + or (self.nccl_gather and self.gather_to_rank0) + ) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (EncoderCausal3D, DecoderCausal3D)): + module.gradient_checkpointing = value + + def enable_temporal_tiling(self, use_tiling: bool = True): + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.enable_spatial_tiling(use_tiling) + self.enable_temporal_tiling(use_tiling) + + def disable_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.disable_spatial_tiling() + self.disable_temporal_tiling() + + def enable_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + + def load_trt_decoder(self): + self.use_trt_decoder = True + self.engine = EngineFromBytes(BytesFromPath(self.engine_path)) + + self.trt_decoder_runner = TrtRunner(self.engine) + self.activate_trt_decoder() + + def disable_trt_decoder(self): + self.use_trt_decoder = False + del self.engine + + def activate_trt_decoder(self): + self.trt_decoder_runner.activate() + + def deactivate_trt_decoder(self): + self.trt_decoder_runner.deactivate() + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + @apply_forward_hook + def encode( + self, x: torch.FloatTensor, return_dict: bool = True + ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + """ + Encode a batch of images into latents. + + Args: + x (`torch.FloatTensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + assert len(x.shape) == 5, "The input tensor should have 5 dimensions" + + if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x, return_dict=return_dict) + + if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size): + return self.spatial_tiled_encode(x, return_dict=return_dict) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self.encoder(x) + + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + assert len(z.shape) == 5, "The input tensor should have 5 dimensions" + + if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z, return_dict=return_dict) + + if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): + return self.spatial_tiled_decode(z, return_dict=return_dict) + + if self.use_trt_decoder: + # For unknown reason, `copy_outputs_to_host` must be set to True + dec = self.trt_decoder_runner.infer({"input": z.to(RECOMMENDED_DTYPE).contiguous()}, copy_outputs_to_host=True)["output"].to(device=z.device, dtype=z.dtype) + else: + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode( + self, z: torch.FloatTensor, return_dict: bool = True, generator=None + ) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Decode a batch of images. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + + if self.parallel_decode: + if z.dtype != RECOMMENDED_DTYPE: + loguru.logger.warning( + f'For better performance, using {RECOMMENDED_DTYPE} for both latent features and model parameters is recommended.' + f'Current latent dtype {z.dtype}. ' + f'Please note that the input latent will be cast to {RECOMMENDED_DTYPE} internally when decoding.' + ) + z = z.to(RECOMMENDED_DTYPE) + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + if blend_extent == 0: + return b + + a_region = a[..., -blend_extent:, :] + b_region = b[..., :blend_extent, :] + + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent + weights = weights.view(1, 1, 1, blend_extent, 1) + + blended = a_region * (1 - weights) + b_region * weights + + b[..., :blend_extent, :] = blended + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + if blend_extent == 0: + return b + + a_region = a[..., -blend_extent:] + b_region = b[..., :blend_extent] + + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent + weights = weights.view(1, 1, 1, 1, blend_extent) + + blended = a_region * (1 - weights) + b_region * weights + + b[..., :blend_extent] = blended + return b + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + if blend_extent == 0: + return b + + a_region = a[..., -blend_extent:, :, :] + b_region = b[..., :blend_extent, :, :] + + weights = torch.arange(blend_extent, device=a.device, dtype=a.dtype) / blend_extent + weights = weights.view(1, 1, blend_extent, 1, 1) + + blended = a_region * (1 - weights) + b_region * weights + + b[..., :blend_extent, :, :] = blended + return b + + def spatial_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True, return_moments: bool = False) -> AutoencoderKLOutput: + r"""Encode a batch of images using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.FloatTensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split video into tiles and encode them separately. + rows = [] + for i in range(0, x.shape[-2], overlap_size): + row = [] + for j in range(0, x.shape[-1], overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + moments = torch.cat(result_rows, dim=-2) + if return_moments: + return moments + + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + + def spatial_tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + if self.parallel_decode: + + rank = mpi_rank() + torch.cuda.set_device(rank) # set device for trt_runner + world_size = mpi_world_size() + + tiles = [] + afters_if_padding = [] + for i in range(0, z.shape[-2], overlap_size): + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + + if self.use_padding and (tile.shape[-2] < self.tile_latent_min_size or tile.shape[-1] < self.tile_latent_min_size): + from torch.nn import functional as F + after_h = tile.shape[-2] * 8 + after_w = tile.shape[-1] * 8 + padding = (0, self.tile_latent_min_size - tile.shape[-1], 0, self.tile_latent_min_size - tile.shape[-2], 0, 0) + tile = F.pad(tile, padding, "replicate").to(device=tile.device, dtype=tile.dtype) + afters_if_padding.append((after_h, after_w)) + else: + afters_if_padding.append(None) + + tiles.append(tile) + + + # balance tasks + ratio = math.ceil(len(tiles) / world_size) + tiles_curr_rank = tiles[rank * ratio: None if rank == world_size - 1 else (rank + 1) * ratio] + + decoded_results = [] + + + total = len(tiles) + n_task = ([ratio] * (total // ratio) + ([total % ratio] if total % ratio else [])) + n_task = n_task + [0] * (8 - len(n_task)) + + for i, tile in enumerate(tiles_curr_rank): + if self.use_trt_decoder: + # For unknown reason, `copy_outputs_to_host` must be set to True + decoded = self.trt_decoder_runner.infer( + {"input": tile.to(RECOMMENDED_DTYPE).contiguous()}, + copy_outputs_to_host=True + )["output"].to(device=z.device, dtype=z.dtype) + decoded_results.append(decoded) + else: + decoded_results.append(self.decoder(self.post_quant_conv(tile))) + + + def find(n): + return next((i for i, task_n in enumerate(n_task) if task_n < n), len(n_task)) + + + if self.nccl_gather and self.gather_to_rank0: + self.igather.gather(decoded, n_rank=find(i + 1)) + + if not self.nccl_gather: + if self.gather_to_rank0: + decoded_results = mpi_comm().gather(decoded_results, root=0) + if rank != 0: + return DecoderOutput(sample=None) + else: + decoded_results = mpi_comm().allgather(decoded_results) + + decoded_results = sum(decoded_results, []) + else: + # [Kevin]: + # We expect all tiles obtained from the same rank have the same shape. + # Shapes among ranks can differ due to the imbalance of task assignment. + if self.gather_to_rank0: + if rank == 0: + self.igather.wait() + gather_results = self.igather.buffers + self.igather.clear() + else: + raise NotImplementedError('The old `allgather` implementation is deprecated for nccl plan.') + + if rank != 0 and self.gather_to_rank0: + return DecoderOutput(sample=None) + + decoded_results = [col[i] for i in range(max([len(k) for k in gather_results])) for col in gather_results if i < len(col)] + + + # Crop the padding region in pixel level + if self.use_padding: + new_decoded_results = [] + for after, dec in zip(afters_if_padding, decoded_results): + if after is not None: + after_h, after_w = after + new_decoded_results.append(dec[:, :, :, :after_h, :after_w]) + else: + new_decoded_results.append(dec) + decoded_results = new_decoded_results + + rows = [] + decoded_results_iter = iter(decoded_results) + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + row.append(next(decoded_results_iter).to(rank)) + rows.append(row) + else: + rows = [] + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: + assert not self.disable_causal_conv, "Temporal tiling is only compatible with causal convolutions." + + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_latent_min_tsize - blend_extent + + # Split the video into tiles and encode them separately. + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size): + tile = self.spatial_tiled_encode(tile, return_moments=True) + else: + tile = self.encoder(tile) + tile = self.quant_conv(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit+1, :, :]) + + moments = torch.cat(result_row, dim=2) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def temporal_tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + # Split z into overlapping tiles and decode them separately. + assert not self.disable_causal_conv, "Temporal tiling is only supported with causal convolutions." + + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - blend_extent + rank = 0 if CPU_OFFLOAD or DISABLE_SP else mpi_rank() + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size): + decoded = self.spatial_tiled_decode(tile, return_dict=True).sample + else: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + if i > 0 and (not (self.parallel_decode and self.gather_to_rank0) or rank == 0): + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + if not CPU_OFFLOAD and not DISABLE_SP and self.parallel_decode and self.gather_to_rank0 and rank != 0: + return DecoderOutput(sample=None) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit+1, :, :]) + + dec = torch.cat(result_row, dim=2) + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.FloatTensor, + sample_posterior: bool = False, + return_dict: bool = True, + return_posterior: bool = False, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput2, torch.FloatTensor]: + r""" + Args: + sample (`torch.FloatTensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z).sample + + if not return_dict: + if return_posterior: + return (dec, posterior) + else: + return (dec,) + if return_posterior: + return DecoderOutput2(sample=dec, posterior=posterior) + else: + return DecoderOutput2(sample=dec) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections + def fuse_qkv_projections(self): + """ + Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, + key, value) are fused. For cross-attention modules, key and value projection matrices are fused. + + + + This API is 🧪 experimental. + + + """ + self.original_attn_processors = None + + for _, attn_processor in self.attn_processors.items(): + if "Added" in str(attn_processor.__class__.__name__): + raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") + + self.original_attn_processors = self.attn_processors + + for module in self.modules(): + if isinstance(module, Attention): + module.fuse_projections(fuse=True) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections + def unfuse_qkv_projections(self): + """Disables the fused QKV projection if enabled. + + + + This API is 🧪 experimental. + + + + """ + if self.original_attn_processors is not None: + self.set_attn_processor(self.original_attn_processors) diff --git a/hymm_sp/vae/unet_causal_3d_blocks.py b/hymm_sp/vae/unet_causal_3d_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..021c7cf21b2da91ded87c1710451b1cf21b47c46 --- /dev/null +++ b/hymm_sp/vae/unet_causal_3d_blocks.py @@ -0,0 +1,884 @@ +# Copyright 2023 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. +from typing import Any, Dict, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from einops import rearrange + +from diffusers.utils import is_torch_version, logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import SpatialNorm +from diffusers.models.attention_processor import Attention +from diffusers.models.normalization import AdaGroupNorm +from diffusers.models.normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None): + seq_len = n_frame * n_hw + mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device) + for i in range(seq_len): + i_frame = i // n_hw + mask[i, : (i_frame + 1) * n_hw] = 0 + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class CausalConv3d(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode = 'replicate', + disable_causal=False, + **kwargs + ): + super().__init__() + + self.pad_mode = pad_mode + if disable_causal: + padding = (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2) + else: + padding = (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size - 1, 0) # W, H, T + self.time_causal_padding = padding + + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride = stride, dilation = dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + +class CausalAvgPool3d(nn.Module): + def __init__( + self, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]], + pad_mode = 'replicate', + disable_causal=False, + **kwargs + ): + super().__init__() + + self.pad_mode = pad_mode + if disable_causal: + padding = (0, 0, 0, 0, 0, 0) + else: + padding = (0, 0, 0, 0, stride - 1, 0) # W, H, T + self.time_causal_padding = padding + + self.conv = nn.AvgPool3d(kernel_size, stride=stride, ceil_mode=True, **kwargs) + self.pad_mode = pad_mode + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + +class UpsampleCausal3D(nn.Module): + """A 3D upsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + use_conv_transpose (`bool`, default `False`): + option to use a convolution transpose. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + name (`str`, default `conv`): + name of the upsampling 3D layer. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + kernel_size: Optional[int] = None, + padding=1, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + interpolate=True, + upsample_factor=(2, 2, 2), + disable_causal=False, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + self.interpolate = interpolate + self.upsample_factor = upsample_factor + self.disable_causal = disable_causal + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + conv = None + if use_conv_transpose: + assert False, "Not Implement yet" + if kernel_size is None: + kernel_size = 4 + conv = nn.ConvTranspose2d( + channels, self.out_channels, kernel_size=kernel_size, stride=2, padding=padding, bias=bias + ) + elif use_conv: + if kernel_size is None: + kernel_size = 3 + conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias, disable_causal=disable_causal) + + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + output_size: Optional[int] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + assert False, "Not Implement yet" + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + if self.use_conv_transpose: + return self.conv(hidden_states) + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + # https://github.com/pytorch/pytorch/issues/86679 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if self.interpolate: + B, C, T, H, W = hidden_states.shape + if not self.disable_causal: + first_h, other_h = hidden_states.split((1, T-1), dim=2) + if output_size is None: + if T > 1: + other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest") + + first_h = first_h.squeeze(2) + first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest") + first_h = first_h.unsqueeze(2) + else: + assert False, "Not Implement yet" + other_h = F.interpolate(other_h, size=output_size, mode="nearest") + + if T > 1: + hidden_states = torch.cat((first_h, other_h), dim=2) + else: + hidden_states = first_h + else: + hidden_states = F.interpolate(hidden_states, scale_factor=self.upsample_factor, mode="nearest") + + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + +class DownsampleCausal3D(nn.Module): + """A 3D downsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + padding (`int`, default `1`): + padding for the convolution. + name (`str`, default `conv`): + name of the downsampling 3D layer. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + kernel_size=3, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + stride=2, + disable_causal=False, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = stride + self.name = name + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + if use_conv: + conv = CausalConv3d( + self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, disable_causal=disable_causal, bias=bias + ) + else: + raise NotImplementedError + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states) + + return hidden_states + +class ResnetBlockCausal3D(nn.Module): + r""" + A Resnet block. + + Parameters: + in_channels (`int`): The number of channels in the input. + out_channels (`int`, *optional*, default to be `None`): + The number of output channels for the first conv2d layer. If None, same as `in_channels`. + dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. + temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. + groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. + groups_out (`int`, *optional*, default to None): + The number of groups to use for the second normalization layer. if set to None, same as `groups`. + eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. + non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. + time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config. + By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" or + "ada_group" for a stronger conditioning with scale and shift. + kernel (`torch.FloatTensor`, optional, default to None): FIR filter, see + [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. + output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. + use_in_shortcut (`bool`, *optional*, default to `True`): + If `True`, add a 1x1 nn.conv2d layer for skip-connection. + up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. + down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. + conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the + `conv_shortcut` output. + conv_3d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. + If None, same as `out_channels`. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + time_embedding_norm: str = "default", # default, scale_shift, ada_group, spatial + kernel: Optional[torch.FloatTensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, + conv_shortcut_bias: bool = True, + conv_3d_out_channels: Optional[int] = None, + disable_causal: bool = False, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.up = up + self.down = down + self.output_scale_factor = output_scale_factor + self.time_embedding_norm = time_embedding_norm + self.skip_time_act = skip_time_act + + linear_cls = nn.Linear + + if groups_out is None: + groups_out = groups + + if self.time_embedding_norm == "ada_group": + self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm1 = SpatialNorm(in_channels, temb_channels) + else: + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1, disable_causal=disable_causal) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + self.time_emb_proj = linear_cls(temb_channels, out_channels) + elif self.time_embedding_norm == "scale_shift": + self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels) + elif self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial": + self.time_emb_proj = None + else: + raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ") + else: + self.time_emb_proj = None + + if self.time_embedding_norm == "ada_group": + self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm2 = SpatialNorm(out_channels, temb_channels) + else: + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) + + self.dropout = torch.nn.Dropout(dropout) + conv_3d_out_channels = conv_3d_out_channels or out_channels + self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1, disable_causal=disable_causal) + + self.nonlinearity = get_activation(non_linearity) + + self.upsample = self.downsample = None + if self.up: + self.upsample = UpsampleCausal3D(in_channels, use_conv=False, disable_causal=disable_causal) + elif self.down: + self.downsample = DownsampleCausal3D(in_channels, use_conv=False, disable_causal=disable_causal, name="op") + + self.use_in_shortcut = self.in_channels != conv_3d_out_channels if use_in_shortcut is None else use_in_shortcut + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = CausalConv3d( + in_channels, + conv_3d_out_channels, + kernel_size=1, + stride=1, + disable_causal=disable_causal, + bias=conv_shortcut_bias, + ) + + def forward( + self, + input_tensor: torch.FloatTensor, + temb: torch.FloatTensor, + scale: float = 1.0, + ) -> torch.FloatTensor: + hidden_states = input_tensor + + if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial": + hidden_states = self.norm1(hidden_states, temb) + else: + hidden_states = self.norm1(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = ( + self.upsample(input_tensor, scale=scale) + ) + hidden_states = ( + self.upsample(hidden_states, scale=scale) + ) + elif self.downsample is not None: + input_tensor = ( + self.downsample(input_tensor, scale=scale) + ) + hidden_states = ( + self.downsample(hidden_states, scale=scale) + ) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = ( + self.time_emb_proj(temb, scale)[:, :, None, None] + ) + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial": + hidden_states = self.norm2(hidden_states, temb) + else: + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = ( + self.conv_shortcut(input_tensor) + ) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + +def get_down_block3d( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + downsample_stride: int, + resnet_eps: float, + resnet_act_fn: str, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + downsample_type: Optional[str] = None, + dropout: float = 0.0, + disable_causal: bool = False, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type + if down_block_type == "DownEncoderBlockCausal3D": + return DownEncoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + downsample_stride=downsample_stride, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + disable_causal=disable_causal, + ) + raise ValueError(f"{down_block_type} does not exist.") + +def get_up_block3d( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + upsample_scale_factor: Tuple, + resnet_eps: float, + resnet_act_fn: str, + resolution_idx: Optional[int] = None, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + upsample_type: Optional[str] = None, + dropout: float = 0.0, + disable_causal: bool = False, +) -> nn.Module: + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + if up_block_type == "UpDecoderBlockCausal3D": + return UpDecoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + upsample_scale_factor=upsample_scale_factor, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + disable_causal=disable_causal, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlockCausal3D(nn.Module): + """ + A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks. + + Args: + in_channels (`int`): The number of input channels. + temb_channels (`int`): The number of temporal embedding channels. + dropout (`float`, *optional*, defaults to 0.0): The dropout rate. + num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. + resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. + resnet_time_scale_shift (`str`, *optional*, defaults to `default`): + The type of normalization to apply to the time embeddings. This can help to improve the performance of the + model on tasks with long-range temporal dependencies. + resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. + resnet_groups (`int`, *optional*, defaults to 32): + The number of groups to use in the group normalization layers of the resnet blocks. + attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks. + resnet_pre_norm (`bool`, *optional*, defaults to `True`): + Whether to use pre-normalization for the resnet blocks. + add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. + attention_head_dim (`int`, *optional*, defaults to 1): + Dimension of a single attention head. The number of attention heads is determined based on this value and + the number of input channels. + output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. + + Returns: + `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size, + in_channels, height, width)`. + + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + disable_causal: bool = False, + causal_attention: bool = False, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + self.causal_attention = causal_attention + + if attn_groups is None: + attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None + + # there is always at least one resnet + resnets = [ + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + disable_causal=disable_causal, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + #assert False, "Not implemented yet" + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + disable_causal=disable_causal, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + B, C, T, H, W = hidden_states.shape + hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c") + if self.causal_attention: + attention_mask = prepare_causal_attention_mask(T, H * W, hidden_states.dtype, hidden_states.device, batch_size=B) + else: + attention_mask = None + hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask) + hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class DownEncoderBlockCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + disable_causal: bool = False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + disable_causal=disable_causal, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + DownsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + stride=downsample_stride, + disable_causal=disable_causal, + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class UpDecoderBlockCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + upsample_scale_factor = (2, 2, 2), + temb_channels: Optional[int] = None, + disable_causal: bool = False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + disable_causal=disable_causal, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + UpsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + disable_causal=disable_causal + ) + ] + ) + else: + self.upsamplers = None + + self.resolution_idx = resolution_idx + + def forward( + self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0 + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + diff --git a/hymm_sp/vae/vae.py b/hymm_sp/vae/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..b7198a30bb3b5aaa283579cdf4e287f2906de2e8 --- /dev/null +++ b/hymm_sp/vae/vae.py @@ -0,0 +1,427 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn + +from diffusers.utils import BaseOutput, is_torch_version +from diffusers.utils.torch_utils import randn_tensor +from diffusers.models.attention_processor import SpatialNorm +from .unet_causal_3d_blocks import ( + CausalConv3d, + UNetMidBlockCausal3D, + get_down_block3d, + get_up_block3d, +) + +@dataclass +class DecoderOutput(BaseOutput): + r""" + Output of decoding method. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The decoded output sample from the last layer of the model. + """ + + sample: torch.FloatTensor + + +class EncoderCausal3D(nn.Module): + r""" + The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available + options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + double_z (`bool`, *optional*, defaults to `True`): + Whether to double the number of output channels for the last block. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + disable_causal: bool = False, + mid_block_causal_attn: bool = False, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, disable_causal=disable_causal) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block) + elif time_compression_ratio == 8: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2, ) if add_time_downsample else (1, ) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + down_block = get_down_block3d( + down_block_type, + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + downsample_stride=downsample_stride, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None, + disable_causal=disable_causal, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + disable_causal=disable_causal, + causal_attention=mid_block_causal_attn, + ) + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3, disable_causal=disable_causal) + + self.gradient_checkpointing = False + + def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: + r"""The forward method of the `EncoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions" + + sample = self.conv_in(sample) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + # down + if is_torch_version(">=", "1.11.0"): + for down_block in self.down_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(down_block), sample, use_reentrant=False + ) + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), sample, use_reentrant=False + ) + else: + for down_block in self.down_blocks: + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(down_block), sample) + # middle + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample) + + else: + # down + for down_block in self.down_blocks: + sample = down_block(sample) + + # middle + sample = self.mid_block(sample) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DecoderCausal3D(nn.Module): + r""" + The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + The types of up blocks to use. See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + norm_type (`str`, *optional*, defaults to `"group"`): + The normalization type to use. Can be either `"group"` or `"spatial"`. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + disable_causal: bool = False, + mid_block_causal_attn: bool = False, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, disable_causal=disable_causal) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + disable_causal=disable_causal, + causal_attention=mid_block_causal_attn, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block) + elif time_compression_ratio == 8: + add_spatial_upsample = bool(i >= len(block_out_channels) - num_spatial_upsample_layers) + add_time_upsample = bool(i >= len(block_out_channels) - num_time_upsample_layers) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + up_block = get_up_block3d( + up_block_type, + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + prev_output_channel=None, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=temb_channels, + resnet_time_scale_shift=norm_type, + disable_causal=disable_causal, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) + else: + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3, disable_causal=disable_causal) + + self.gradient_checkpointing = False + + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + r"""The forward method of the `DecoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions" + + sample = self.conv_in(sample) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), + sample, + latent_embeds, + use_reentrant=False, + ) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + use_reentrant=False, + ) + else: + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), sample, latent_embeds + ) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds) + else: + # middle + sample = self.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds) + + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + reduce_dim = list(range(1, self.mean.ndim)) + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=reduce_dim, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=reduce_dim, + ) + + def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..cd728cd8949a93728911e047e5fade18185c8c0d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,19 @@ +opencv-python==4.9.0.80 +diffusers==0.33.0 +transformers==4.45.1 +accelerate==1.1.1 +pandas==2.0.3 +numpy==1.24.4 +einops==0.7.0 +tqdm==4.66.2 +loguru==0.7.2 +imageio==2.34.0 +imageio-ffmpeg==0.5.1 +safetensors==0.4.3 +gradio==4.42.0 +fastapi==0.115.12 +uvicorn==0.34.2 +decord==0.6.0 +librosa==0.11.0 +scikit-video==1.1.11 +ffmpeg \ No newline at end of file diff --git a/scripts/run_gradio.sh b/scripts/run_gradio.sh new file mode 100644 index 0000000000000000000000000000000000000000..58d86c7b79d35f65c0a547b35a0efe8dfef52b28 --- /dev/null +++ b/scripts/run_gradio.sh @@ -0,0 +1,22 @@ +#!/bin/bash +JOBS_DIR=$(dirname $(dirname "$0")) +export PYTHONPATH=./ + +export MODEL_BASE=./weights + +checkpoint_path=${MODEL_BASE}/ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt + + +torchrun --nnodes=1 --nproc_per_node=8 --master_port 29605 hymm_gradio/flask_audio.py \ + --input 'assets/test.csv' \ + --ckpt ${checkpoint_path} \ + --sample-n-frames 129 \ + --seed 128 \ + --image-size 704 \ + --cfg-scale 7.5 \ + --infer-steps 50 \ + --use-deepcache 1 \ + --flow-shift-eval-video 5.0 & + + +python3 hymm_gradio/gradio_audio.py diff --git a/scripts/run_sample_batch_sp.sh b/scripts/run_sample_batch_sp.sh new file mode 100644 index 0000000000000000000000000000000000000000..72fe93c0eb1d57ab77ca9a358c1dd2aafaf50c41 --- /dev/null +++ b/scripts/run_sample_batch_sp.sh @@ -0,0 +1,20 @@ +#!/bin/bash +JOBS_DIR=$(dirname $(dirname "$0")) +export PYTHONPATH=./ + +export MODEL_BASE=./weights +OUTPUT_BASEPATH=./results +checkpoint_path=${MODEL_BASE}/ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt + + +torchrun --nnodes=1 --nproc_per_node=8 --master_port 29605 hymm_sp/sample_batch.py \ + --input 'assets/test.csv' \ + --ckpt ${checkpoint_path} \ + --sample-n-frames 129 \ + --seed 128 \ + --image-size 704 \ + --cfg-scale 7.5 \ + --infer-steps 50 \ + --use-deepcache 1 \ + --flow-shift-eval-video 5.0 \ + --save-path ${OUTPUT_BASEPATH} diff --git a/scripts/run_single_audio.sh b/scripts/run_single_audio.sh new file mode 100644 index 0000000000000000000000000000000000000000..08f8273d21427f29afbc67162ca71e3361678d41 --- /dev/null +++ b/scripts/run_single_audio.sh @@ -0,0 +1,25 @@ +#!/bin/bash +JOBS_DIR=$(dirname $(dirname "$0")) +export PYTHONPATH=./ + +export MODEL_BASE=./weights +OUTPUT_BASEPATH=./results-single + +# checkpoint_path=${MODEL_BASE}/ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt +checkpoint_path=${MODEL_BASE}/ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt + + +export DISABLE_SP=1 +CUDA_VISIBLE_DEVICES=0 python3 hymm_sp/sample_gpu_poor.py \ + --input 'assets/test.csv' \ + --ckpt ${checkpoint_path} \ + --sample-n-frames 129 \ + --seed 128 \ + --image-size 704 \ + --cfg-scale 7.5 \ + --infer-steps 50 \ + --use-deepcache 1 \ + --flow-shift-eval-video 5.0 \ + --save-path ${OUTPUT_BASEPATH} \ + --use-fp8 \ + --infer-min diff --git a/scripts/run_single_poor.sh b/scripts/run_single_poor.sh new file mode 100644 index 0000000000000000000000000000000000000000..ef94046e7c755f8bba4205ba81b82ce2f41e237d --- /dev/null +++ b/scripts/run_single_poor.sh @@ -0,0 +1,24 @@ +#!/bin/bash +JOBS_DIR=$(dirname $(dirname "$0")) +export PYTHONPATH=./ + +export MODEL_BASE=./weights +OUTPUT_BASEPATH=./results-poor + +checkpoint_path=${MODEL_BASE}/ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt + +export CPU_OFFLOAD=1 +CUDA_VISIBLE_DEVICES=0 python3 hymm_sp/sample_gpu_poor.py \ + --input 'assets/test.csv' \ + --ckpt ${checkpoint_path} \ + --sample-n-frames 129 \ + --seed 128 \ + --image-size 704 \ + --cfg-scale 7.5 \ + --infer-steps 50 \ + --use-deepcache 1 \ + --flow-shift-eval-video 5.0 \ + --save-path ${OUTPUT_BASEPATH} \ + --use-fp8 \ + --cpu-offload \ + --infer-min diff --git a/weights/README.md b/weights/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5864c2f44ef6313636db0027b854db029cc39467 --- /dev/null +++ b/weights/README.md @@ -0,0 +1,44 @@ +# Download Pretrained Models + +All models are stored in `HunyuanVideo-Avatar/weights` by default, and the file structure is as follows +```shell +HunyuanVideo-Avatar + ├──weights + │ ├──ckpts + │ │ ├──README.md + │ │ ├──hunyuan-video-t2v-720p + │ │ │ ├──transformers + │ │ │ │ ├──mp_rank_00_model_states.pt + │ │ │ │ ├──mp_rank_00_model_states_fp8.pt + │ │ │ │ ├──mp_rank_00_model_states_fp8_map.pt + │ │ │ ├──vae + │ │ │ │ ├──pytorch_model.pt + │ │ │ │ ├──config.json + │ │ ├──llava_llama_image + │ │ │ ├──model-00001-of-00004.safatensors + │ │ │ ├──model-00002-of-00004.safatensors + │ │ │ ├──model-00003-of-00004.safatensors + │ │ │ ├──model-00004-of-00004.safatensors + │ │ │ ├──... + │ │ ├──text_encoder_2 + │ │ ├──whisper-tiny + │ │ ├──det_align + │ │ ├──... +``` + +## Download HunyuanVideo-Avatar model +To download the HunyuanCustom model, first install the huggingface-cli. (Detailed instructions are available [here](https://huggingface.co/docs/huggingface_hub/guides/cli).) + +```shell +python -m pip install "huggingface_hub[cli]" +``` + +Then download the model using the following commands: + +```shell +# Switch to the directory named 'HunyuanVideo-Avatar/weights' +cd HunyuanVideo-Avatar/weights +# Use the huggingface-cli tool to download HunyuanVideo-Avatar model in HunyuanVideo-Avatar/weights dir. +# The download time may vary from 10 minutes to 1 hour depending on network conditions. +huggingface-cli download tencent/HunyuanVideo-Avatar --local-dir ./ +```