File size: 5,308 Bytes
226c7c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 cosmos_transfer1.diffusion.config.transfer.blurs import BlurAugmentorConfig, random_blur_config
from cosmos_transfer1.diffusion.config.transfer.conditioner import CTRL_AUG_KEYS, CTRL_HINT_KEYS, CTRL_HINT_KEYS_COMB
from cosmos_transfer1.diffusion.datasets.augmentors.basic_augmentors import (
ReflectionPadding,
ResizeLargestSideAspectPreserving,
)
from cosmos_transfer1.diffusion.datasets.augmentors.control_input import (
VIDEO_RES_SIZE_INFO,
AddControlInput,
AddControlInputComb,
)
from cosmos_transfer1.diffusion.datasets.augmentors.merge_datadict import DataDictMerger
from cosmos_transfer1.utils.lazy_config import LazyCall as L
AUGMENTOR_OPTIONS = {}
def augmentor_register(key):
def decorator(func):
AUGMENTOR_OPTIONS[key] = func
return func
return decorator
@augmentor_register("video_basic_augmentor")
def get_video_augmentor(
resolution: str,
blur_config=None,
):
return {
"merge_datadict": L(DataDictMerger)(
input_keys=["video"],
output_keys=[
"video",
"fps",
"num_frames",
"frame_start",
"frame_end",
"orig_num_frames",
],
),
"resize_largest_side_aspect_ratio_preserving": L(ResizeLargestSideAspectPreserving)(
input_keys=["video"],
args={"size": VIDEO_RES_SIZE_INFO[resolution]},
),
"reflection_padding": L(ReflectionPadding)(
input_keys=["video"],
args={"size": VIDEO_RES_SIZE_INFO[resolution]},
),
}
"""
register all the video ctrlnet augmentors for data loading
"""
for hint_key in CTRL_HINT_KEYS:
def get_video_ctrlnet_augmentor(hint_key, use_random=True):
def _get_video_ctrlnet_augmentor(
resolution: str,
blur_config: BlurAugmentorConfig = random_blur_config,
):
if hint_key == "control_input_keypoint":
add_control_input = L(AddControlInputComb)(
input_keys=["", "video"],
output_keys=[hint_key],
args={
"comb": CTRL_HINT_KEYS_COMB[hint_key],
"use_openpose_format": True,
"kpt_thr": 0.6,
"human_kpt_line_width": 4,
},
use_random=use_random,
blur_config=blur_config,
)
elif hint_key in CTRL_HINT_KEYS_COMB:
add_control_input = L(AddControlInputComb)(
input_keys=["", "video"],
output_keys=[hint_key],
args={"comb": CTRL_HINT_KEYS_COMB[hint_key]},
use_random=use_random,
blur_config=blur_config,
)
else:
add_control_input = L(AddControlInput)(
input_keys=["", "video"],
output_keys=[hint_key],
use_random=use_random,
blur_config=blur_config,
)
input_keys = ["video"]
output_keys = [
"video",
"fps",
"num_frames",
"frame_start",
"frame_end",
"orig_num_frames",
]
for key, value in CTRL_AUG_KEYS.items():
if key in hint_key:
input_keys.append(value)
output_keys.append(value)
augmentation = {
# "merge_datadict": L(DataDictMerger)(
# input_keys=input_keys,
# output_keys=output_keys,
# ),
# this addes the control input tensor to the data dict
"add_control_input": add_control_input,
# this resizes both the video and the control input to the model's required input size
"resize_largest_side_aspect_ratio_preserving": L(ResizeLargestSideAspectPreserving)(
input_keys=["video", hint_key],
args={"size": VIDEO_RES_SIZE_INFO[resolution]},
),
"reflection_padding": L(ReflectionPadding)(
input_keys=["video", hint_key],
args={"size": VIDEO_RES_SIZE_INFO[resolution]},
),
}
return augmentation
return _get_video_ctrlnet_augmentor
augmentor_register(f"video_ctrlnet_augmentor_{hint_key}")(get_video_ctrlnet_augmentor(hint_key))
|